Refactored makefile, moved around files to more appropriate places, deleted unused files,

reformatted code.

Makefile

@@ -1,18 +1,9 @@
 setup:
-	conda activate manim
+	conda activate manim; \
 	export PROJECT_ROOT=$(pwd)
-video:
-	manim -pqh src/variational_autoencoder.py VAEScene --media_dir media
-	cp media/videos/vae/720p60/VAEScene.mp4 examples
-train:
-	cd src/autoencoder_models
-	python vanilla_autoencoder.py
-	python variational_autoencoder.py
-generate_visualizations:
-	cd src/autoencoder_models
-	python generate_images.py
-	python generate_interpolation.py
-	python generate_disentanglement.py
 checkstyle:
-	pycodestyle src
+	black .; \
-	pydocstyle src
+	pydocstyle .
+publish_pip:
+	python3 -m build; \
+	python3 -m twine upload --repository pypi dist/*

examples/basic_neural_network/basic_neural_network.py

@@ -8,11 +8,7 @@ class NeuralNetworkScene(Scene):
 
     def construct(self):
         # Make the Layer object
-        layers = [
+        layers = [FeedForwardLayer(3), FeedForwardLayer(5), FeedForwardLayer(3)]
-            FeedForwardLayer(3), 
-            FeedForwardLayer(5), 
-            FeedForwardLayer(3)
-        ]
         nn = NeuralNetwork(layers)
         nn.scale(2)
         nn.move_to(ORIGIN)

examples/cnn/activation_functions.py

@@ -15,6 +15,7 @@ config.frame_height = 7.0
 config.frame_width = 7.0
 ROOT_DIR = Path(__file__).parents[2]
 
+
 def make_code_snippet():
     code_str = """
         # Make the neural network
@@ -42,6 +43,7 @@ def make_code_snippet():
 
     return code
 
+
 class CombinedScene(ThreeDScene):
     def construct(self):
         image = Image.open(ROOT_DIR / "assets/mnist/digit.jpeg")
@@ -72,4 +74,4 @@ class CombinedScene(ThreeDScene):
         # Play animation
         forward_pass = nn.make_forward_pass_animation()
         self.wait(1)
-        self.play(forward_pass)
+        self.play(forward_pass)

examples/cnn/cnn.py

@@ -15,6 +15,7 @@ config.frame_height = 7.0
 config.frame_width = 7.0
 ROOT_DIR = Path(__file__).parents[2]
 
+
 def make_code_snippet():
     code_str = """
         # Make nn
@@ -44,6 +45,7 @@ def make_code_snippet():
 
     return code
 
+
 class CombinedScene(ThreeDScene):
     def construct(self):
         image = Image.open(ROOT_DIR / "assets/mnist/digit.jpeg")

examples/cnn/cnn_max_pool.py

@@ -14,6 +14,7 @@ config.pixel_width = 1900
 config.frame_height = 6.0
 config.frame_width = 6.0
 
+
 def make_code_snippet():
     code_str = """
         # Make the neural network
@@ -42,12 +43,14 @@ def make_code_snippet():
 
     return code
 
+
 class CombinedScene(ThreeDScene):
     def construct(self):
         image = Image.open("../../assets/mnist/digit.jpeg")
         numpy_image = np.asarray(image)
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 ImageLayer(numpy_image, height=1.5),
                 Convolutional2DLayer(1, 8, filter_spacing=0.32),
                 MaxPooling2DLayer(kernel_size=2),
@@ -67,4 +70,4 @@ class CombinedScene(ThreeDScene):
         # Play animation
         forward_pass = nn.make_forward_pass_animation()
         self.wait(1)
-        self.play(forward_pass)
+        self.play(forward_pass)

examples/cnn/one_by_one_convolution.py

@@ -15,6 +15,7 @@ config.frame_height = 7.0
 config.frame_width = 7.0
 ROOT_DIR = Path(__file__).parents[2]
 
+
 def make_code_snippet():
     code_str = """
         # Make nn

examples/cnn/padding_example.py

@@ -15,6 +15,7 @@ config.pixel_width = 1900
 config.frame_height = 6.0
 config.frame_width = 6.0
 
+
 def make_code_snippet():
     code_str = """
         # Make nn
@@ -43,27 +44,28 @@ def make_code_snippet():
 
     return code
 
-class CombinedScene(ThreeDScene):
 
+class CombinedScene(ThreeDScene):
     def construct(self):
         # Make nn
         image = Image.open(ROOT_DIR / "assets/mnist/digit.jpeg")
         numpy_image = np.asarray(image)
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 ImageLayer(numpy_image, height=1.5),
                 Convolutional2DLayer(
-                    num_feature_maps=1, 
+                    num_feature_maps=1,
-                    feature_map_size=6, 
+                    feature_map_size=6,
                     padding=1,
-                    padding_dashed=True
+                    padding_dashed=True,
                 ),
                 Convolutional2DLayer(
-                    num_feature_maps=3, 
+                    num_feature_maps=3,
                     feature_map_size=6,
                     filter_size=3,
                     padding=0,
-                    padding_dashed=False
+                    padding_dashed=False,
                 ),
                 FeedForwardLayer(3),
                 FeedForwardLayer(1),

examples/cnn/resnet_block.py

@@ -9,6 +9,7 @@ config.pixel_width = 1900
 config.frame_height = 6.0
 config.frame_width = 6.0
 
+
 def make_code_snippet():
     code_str = """
         # Make the neural network
@@ -37,13 +38,14 @@ def make_code_snippet():
 
     return code
 
-class ConvScene(ThreeDScene):
 
+class ConvScene(ThreeDScene):
     def construct(self):
         image = Image.open("../../assets/mnist/digit.jpeg")
         numpy_image = np.asarray(image)
 
-        nn = NeuralNetwork({
+        nn = NeuralNetwork(
+            {
                 "layer1": Convolutional2DLayer(1, 5, padding=1),
                 "layer2": Convolutional2DLayer(1, 5, 3, padding=1),
                 "layer3": Convolutional2DLayer(1, 5, 3, padding=1),
@@ -60,7 +62,4 @@ class ConvScene(ThreeDScene):
         self.add(code)
         Group(code, nn).move_to(ORIGIN)
 
-        self.play(
+        self.play(nn.make_forward_pass_animation(), run_time=8)
-            nn.make_forward_pass_animation(),
-            run_time=8
-        )

examples/lenet/lenet.py

@@ -18,12 +18,14 @@ config.frame_width = 20.0
 
 ROOT_DIR = Path(__file__).parents[2]
 
+
 class CombinedScene(ThreeDScene):
     def construct(self):
         image = Image.open(ROOT_DIR / "assets/mnist/digit.jpeg")
         numpy_image = np.asarray(image)
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 ImageLayer(numpy_image, height=4.5),
                 Convolutional2DLayer(1, 28),
                 Convolutional2DLayer(6, 28, 5),

examples/paper_visualizations/oracle_guidance/oracle_guidance.py

@@ -13,7 +13,7 @@ from manim_ml.neural_network.layers import FeedForwardLayer, EmbeddingLayer
 from manim_ml.neural_network.layers.util import get_connective_layer
 import os
 
-from manim_ml.probability import GaussianDistribution
+from manim_ml.utils.mobjects.probability import GaussianDistribution
 
 # Make the specific scene
 config.pixel_height = 1200

examples/readme_example/a_simple_feed_forward_network.py

@@ -8,14 +8,17 @@ config.pixel_width = 1200
 config.frame_height = 4.0
 config.frame_width = 4.0
 
+
 class CombinedScene(ThreeDScene):
     def construct(self):
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
-            FeedForwardLayer(num_nodes=3),
+            [
-            FeedForwardLayer(num_nodes=5),
+                FeedForwardLayer(num_nodes=3),
-            FeedForwardLayer(num_nodes=3)
+                FeedForwardLayer(num_nodes=5),
-        ])
+                FeedForwardLayer(num_nodes=3),
+            ]
+        )
         self.add(nn)
         # Center the nn
         nn.move_to(ORIGIN)

examples/readme_example/activation_functions.py

@@ -10,12 +10,16 @@ config.pixel_width = 1900
 config.frame_height = 6.0
 config.frame_width = 6.0
 
+
 class CombinedScene(ThreeDScene):
     def construct(self):
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 Convolutional2DLayer(1, 7, filter_spacing=0.32),
-                Convolutional2DLayer(3, 5, 3, filter_spacing=0.32, activation_function="ReLU"),
+                Convolutional2DLayer(
+                    3, 5, 3, filter_spacing=0.32, activation_function="ReLU"
+                ),
                 FeedForwardLayer(3, activation_function="Sigmoid"),
             ],
             layer_spacing=0.25,
@@ -26,4 +30,4 @@ class CombinedScene(ThreeDScene):
         # Play animation
         forward_pass = nn.make_forward_pass_animation()
         self.play(ChangeSpeed(forward_pass, speedinfo={}), run_time=10)
-        self.wait(1)
+        self.wait(1)

examples/readme_example/animating_the_forward_pass.py

@@ -1,6 +1,10 @@
 from manim import *
 
-from manim_ml.neural_network import Convolutional2DLayer, FeedForwardLayer, NeuralNetwork
+from manim_ml.neural_network import (
+    Convolutional2DLayer,
+    FeedForwardLayer,
+    NeuralNetwork,
+)
 
 # Make the specific scene
 config.pixel_height = 700
@@ -8,10 +12,12 @@ config.pixel_width = 1900
 config.frame_height = 7.0
 config.frame_width = 7.0
 
+
 class CombinedScene(ThreeDScene):
     def construct(self):
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 Convolutional2DLayer(1, 7, 3, filter_spacing=0.32),
                 Convolutional2DLayer(3, 5, 3, filter_spacing=0.32),
                 Convolutional2DLayer(5, 3, 3, filter_spacing=0.18),

examples/readme_example/convolutional_neural_network_with_images.py

@@ -2,7 +2,12 @@ from manim import *
 from PIL import Image
 import numpy as np
 
-from manim_ml.neural_network import Convolutional2DLayer, FeedForwardLayer, NeuralNetwork, ImageLayer
+from manim_ml.neural_network import (
+    Convolutional2DLayer,
+    FeedForwardLayer,
+    NeuralNetwork,
+    ImageLayer,
+)
 
 # Make the specific scene
 config.pixel_height = 700
@@ -10,13 +15,15 @@ config.pixel_width = 1900
 config.frame_height = 7.0
 config.frame_width = 7.0
 
+
 class CombinedScene(ThreeDScene):
     def construct(self):
         # Make nn
         image = Image.open("../../assets/mnist/digit.jpeg")
         numpy_image = np.asarray(image)
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 ImageLayer(numpy_image, height=1.5),
                 Convolutional2DLayer(1, 7, filter_spacing=0.32),
                 Convolutional2DLayer(3, 5, 3, filter_spacing=0.32),

examples/readme_example/convolutional_neural_networks.py

@@ -1,6 +1,10 @@
 from manim import *
 
-from manim_ml.neural_network import Convolutional2DLayer, FeedForwardLayer, NeuralNetwork
+from manim_ml.neural_network import (
+    Convolutional2DLayer,
+    FeedForwardLayer,
+    NeuralNetwork,
+)
 
 # Make the specific scene
 config.pixel_height = 700
@@ -8,10 +12,12 @@ config.pixel_width = 1900
 config.frame_height = 7.0
 config.frame_width = 7.0
 
+
 class CombinedScene(ThreeDScene):
     def construct(self):
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 Convolutional2DLayer(1, 7, 3, filter_spacing=0.32),
                 Convolutional2DLayer(3, 5, 3, filter_spacing=0.32),
                 Convolutional2DLayer(5, 3, 3, filter_spacing=0.18),

examples/readme_example/example.py

@@ -1,6 +1,10 @@
 from manim import *
 
-from manim_ml.neural_network import Convolutional2DLayer, FeedForwardLayer, NeuralNetwork
+from manim_ml.neural_network import (
+    Convolutional2DLayer,
+    FeedForwardLayer,
+    NeuralNetwork,
+)
 
 # Make the specific scene
 config.pixel_height = 700
@@ -8,10 +12,12 @@ config.pixel_width = 1900
 config.frame_height = 7.0
 config.frame_width = 7.0
 
+
 class CombinedScene(ThreeDScene):
     def construct(self):
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 Convolutional2DLayer(1, 7, 3, filter_spacing=0.32),
                 Convolutional2DLayer(3, 5, 3, filter_spacing=0.32),
                 Convolutional2DLayer(5, 3, 3, filter_spacing=0.18),

examples/readme_example/first_neural_network.py

@@ -1,6 +1,10 @@
 from manim import *
 
-from manim_ml.neural_network import Convolutional2DLayer, FeedForwardLayer, NeuralNetwork
+from manim_ml.neural_network import (
+    Convolutional2DLayer,
+    FeedForwardLayer,
+    NeuralNetwork,
+)
 
 # Make the specific scene
 config.pixel_height = 700
@@ -8,10 +12,12 @@ config.pixel_width = 1900
 config.frame_height = 7.0
 config.frame_width = 7.0
 
+
 class CombinedScene(ThreeDScene):
     def construct(self):
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 Convolutional2DLayer(1, 7, 3, filter_spacing=0.32),
                 Convolutional2DLayer(3, 5, 3, filter_spacing=0.32),
                 Convolutional2DLayer(5, 3, 3, filter_spacing=0.18),

examples/readme_example/max_pooling.py

@@ -12,11 +12,12 @@ config.pixel_width = 1900
 config.frame_height = 6.0
 config.frame_width = 6.0
 
-class MaxPoolingScene(ThreeDScene):
 
+class MaxPoolingScene(ThreeDScene):
     def construct(self):
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 Convolutional2DLayer(1, 8),
                 Convolutional2DLayer(3, 6, 3),
                 MaxPooling2DLayer(kernel_size=2),
@@ -30,4 +31,4 @@ class MaxPoolingScene(ThreeDScene):
         # Play animation
         forward_pass = nn.make_forward_pass_animation()
         self.play(ChangeSpeed(forward_pass, speedinfo={}), run_time=10)
-        self.wait(1)
+        self.wait(1)

examples/readme_example/neural_network_dropout.py

@@ -9,11 +9,12 @@ config.pixel_width = 1900
 config.frame_height = 5.0
 config.frame_width = 5.0
 
+
 class DropoutNeuralNetworkScene(Scene):
-    
     def construct(self):
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 FeedForwardLayer(3, rectangle_color=BLUE),
                 FeedForwardLayer(5, rectangle_color=BLUE),
                 FeedForwardLayer(3, rectangle_color=BLUE),
@@ -31,4 +32,4 @@ class DropoutNeuralNetworkScene(Scene):
                 nn, dropout_rate=0.25, do_forward_pass=True
             )
         )
-        self.wait(1)
+        self.wait(1)

examples/readme_example/old_example.py

@@ -1,4 +1,4 @@
-from manim import * 
+from manim import *
 from PIL import Image
 
 from manim_ml.neural_network.layers.convolutional_2d import Convolutional2DLayer
@@ -6,11 +6,12 @@ from manim_ml.neural_network.layers.feed_forward import FeedForwardLayer
 from manim_ml.neural_network.layers.image import ImageLayer
 from manim_ml.neural_network.neural_network import NeuralNetwork
 
-class ConvolutionalNetworkScene(Scene):
 
+class ConvolutionalNetworkScene(Scene):
     def construct(self):
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 Convolutional2DLayer(1, 7, 3, filter_spacing=0.32),
                 Convolutional2DLayer(3, 5, 3, filter_spacing=0.32),
                 Convolutional2DLayer(5, 3, 3, filter_spacing=0.18),
@@ -22,4 +23,4 @@ class ConvolutionalNetworkScene(Scene):
         # Center the nn
         nn.move_to(ORIGIN)
         self.add(nn)
-        self.play(nn.make_forward_pass_animation())
+        self.play(nn.make_forward_pass_animation())

examples/readme_example/setting_up_a_scene.py

@@ -1,9 +1,10 @@
 from manim import *
+
 # Import modules here
 
-class BasicScene(ThreeDScene):
 
+class BasicScene(ThreeDScene):
     def construct(self):
         # Your code goes here
         text = Text("Your first scene!")
-        self.add(text)
+        self.add(text)

manim_ml/diffusion/mcmc.py

@@ -7,7 +7,7 @@ import scipy
 import scipy.stats
 from tqdm import tqdm
 
-from manim_ml.probability import GaussianDistribution
+from manim_ml.utils.mobjects.probability import GaussianDistribution
 
 
 def gaussian_proposal(x, sigma=0.2):

manim_ml/flow/flow.py

@@ -1,22 +0,0 @@
-"""
-    Animated flow charts. 
-"""
-from manim import *
-
-
-class FlowGraph(VGroup):
-    """Graph container"""
-
-    pass
-
-
-class FlowNode(VGroup):
-    """Node in the FlowGraph"""
-
-    pass
-
-
-class DataNode(FlowNode):
-    """Node that outputs data"""
-
-    pass

manim_ml/lazy_animation.py

@@ -1,13 +0,0 @@
-from manim import *
-
-
-class LazyAnimation(Animation):
-    def __init__(self, animation_function):
-        self.animation_function = animation_function
-        super.__init__()
-
-    def begin(self):
-        update_func_anim = UpdateFromFunc(self.neural_network, create_new_connective)
-        self.add
-
-        super.begin()

manim_ml/manifold.py

@@ -1,3 +0,0 @@
-"""
-    Visaulization of a latent Manifold
-"""

manim_ml/neural_network/__init__.py

@@ -1,23 +1,41 @@
 from manim_ml.neural_network.neural_network import NeuralNetwork
 from manim_ml.neural_network.layers.feed_forward import FeedForwardLayer
-from manim_ml.neural_network.layers.convolutional_2d_to_convolutional_2d import Convolutional2DToConvolutional2D
+from manim_ml.neural_network.layers.convolutional_2d_to_convolutional_2d import (
-from manim_ml.neural_network.layers.convolutional_2d_to_feed_forward import Convolutional2DToFeedForward
+    Convolutional2DToConvolutional2D,
-from manim_ml.neural_network.layers.convolutional_2d_to_max_pooling_2d import Convolutional2DToMaxPooling2D
+)
+from manim_ml.neural_network.layers.convolutional_2d_to_feed_forward import (
+    Convolutional2DToFeedForward,
+)
+from manim_ml.neural_network.layers.convolutional_2d_to_max_pooling_2d import (
+    Convolutional2DToMaxPooling2D,
+)
 from manim_ml.neural_network.layers.convolutional_2d import Convolutional2DLayer
-from manim_ml.neural_network.layers.embedding_to_feed_forward import EmbeddingToFeedForward
+from manim_ml.neural_network.layers.embedding_to_feed_forward import (
+    EmbeddingToFeedForward,
+)
 from manim_ml.neural_network.layers.embedding import EmbeddingLayer
-from manim_ml.neural_network.layers.feed_forward_to_embedding import FeedForwardToEmbedding
+from manim_ml.neural_network.layers.feed_forward_to_embedding import (
-from manim_ml.neural_network.layers.feed_forward_to_feed_forward import FeedForwardToFeedForward
+    FeedForwardToEmbedding,
+)
+from manim_ml.neural_network.layers.feed_forward_to_feed_forward import (
+    FeedForwardToFeedForward,
+)
 from manim_ml.neural_network.layers.feed_forward_to_image import FeedForwardToImage
 from manim_ml.neural_network.layers.feed_forward_to_vector import FeedForwardToVector
 from manim_ml.neural_network.layers.feed_forward import FeedForwardLayer
-from manim_ml.neural_network.layers.image_to_convolutional_2d import ImageToConvolutional2DLayer
+from manim_ml.neural_network.layers.image_to_convolutional_2d import (
+    ImageToConvolutional2DLayer,
+)
 from manim_ml.neural_network.layers.image_to_feed_forward import ImageToFeedForward
 from manim_ml.neural_network.layers.image import ImageLayer
-from manim_ml.neural_network.layers.max_pooling_2d_to_convolutional_2d import MaxPooling2DToConvolutional2D
+from manim_ml.neural_network.layers.max_pooling_2d_to_convolutional_2d import (
+    MaxPooling2DToConvolutional2D,
+)
 from manim_ml.neural_network.layers.max_pooling_2d import MaxPooling2DLayer
-from manim_ml.neural_network.layers.paired_query_to_feed_forward import PairedQueryToFeedForward
+from manim_ml.neural_network.layers.paired_query_to_feed_forward import (
+    PairedQueryToFeedForward,
+)
 from manim_ml.neural_network.layers.paired_query import PairedQueryLayer
 from manim_ml.neural_network.layers.triplet_to_feed_forward import TripletToFeedForward
 from manim_ml.neural_network.layers.triplet import TripletLayer
-from manim_ml.neural_network.layers.vector import VectorLayer
+from manim_ml.neural_network.layers.vector import VectorLayer

manim_ml/neural_network/activation_functions/__init__.py

@@ -1,13 +1,12 @@
 from manim_ml.neural_network.activation_functions.relu import ReLUFunction
 from manim_ml.neural_network.activation_functions.sigmoid import SigmoidFunction
 
-name_to_activation_function_map = {
+name_to_activation_function_map = {"ReLU": ReLUFunction, "Sigmoid": SigmoidFunction}
-    "ReLU": ReLUFunction,
+
-    "Sigmoid": SigmoidFunction
-}
 
 def get_activation_function_by_name(name):
-    assert name in name_to_activation_function_map.keys(), \
+    assert (
-        f"Unrecognized activation function {name}"
+        name in name_to_activation_function_map.keys()
+    ), f"Unrecognized activation function {name}"
 
     return name_to_activation_function_map[name]

manim_ml/neural_network/activation_functions/activation_function.py

@@ -4,6 +4,7 @@ import random
 
 import manim_ml.neural_network.activation_functions.relu as relu
 
+
 class ActivationFunction(ABC, VGroup):
     """Abstract parent class for defining activation functions"""
 

manim_ml/neural_network/activation_functions/sigmoid.py

@@ -5,6 +5,7 @@ from manim_ml.neural_network.activation_functions.activation_function import (
     ActivationFunction,
 )
 
+
 class SigmoidFunction(ActivationFunction):
     """Sigmoid Activation Function"""
 

manim_ml/neural_network/architectures/feed_forward.py

@@ -1,5 +1,6 @@
 from manim_ml.neural_network.layers.feed_forward import FeedForwardLayer
 
+
 class FeedForwardNeuralNetwork(NeuralNetwork):
     """NeuralNetwork with just feed forward layers"""
 

manim_ml/neural_network/layers/__init__.py

@@ -10,7 +10,9 @@ from manim_ml.neural_network.layers.image_to_convolutional_2d import (
 from manim_ml.neural_network.layers.max_pooling_2d_to_convolutional_2d import (
     MaxPooling2DToConvolutional2D,
 )
-from manim_ml.neural_network.layers.max_pooling_2d_to_feed_forward import MaxPooling2DToFeedForward
+from manim_ml.neural_network.layers.max_pooling_2d_to_feed_forward import (
+    MaxPooling2DToFeedForward,
+)
 from .convolutional_2d_to_convolutional_2d import Convolutional2DToConvolutional2D
 from .convolutional_2d import Convolutional2DLayer
 from .feed_forward_to_vector import FeedForwardToVector

manim_ml/neural_network/layers/convolutional_2d.py

@@ -10,7 +10,8 @@ from manim_ml.neural_network.layers.parent_layers import (
     ThreeDLayer,
     VGroupNeuralNetworkLayer,
 )
-from manim_ml.gridded_rectangle import GriddedRectangle
+from manim_ml.utils.mobjects.gridded_rectangle import GriddedRectangle
+
 
 class FeatureMap(VGroup):
     """Class for making a feature map"""
@@ -25,7 +26,7 @@ class FeatureMap(VGroup):
         padding=(0, 0),
         stroke_width=2.0,
         show_grid_lines=False,
-        padding_dashed=False
+        padding_dashed=False,
     ):
         super().__init__()
         self.color = color
@@ -40,8 +41,12 @@ class FeatureMap(VGroup):
         # Check if we have non-zero padding
         if padding[0] > 0 or padding[1] > 0:
             # Make the exterior rectangle dashed
-            width_with_padding = (self.feature_map_size[0] + self.padding[0] * 2) * self.cell_width
+            width_with_padding = (
-            height_with_padding = (self.feature_map_size[1] + self.padding[1] * 2) * self.cell_width
+                self.feature_map_size[0] + self.padding[0] * 2
+            ) * self.cell_width
+            height_with_padding = (
+                self.feature_map_size[1] + self.padding[1] * 2
+            ) * self.cell_width
             self.untransformed_width = width_with_padding
             self.untransformed_height = height_with_padding
 
@@ -58,7 +63,7 @@ class FeatureMap(VGroup):
                 grid_stroke_width=self.stroke_width / 2,
                 grid_stroke_color=self.color,
                 show_grid_lines=self.show_grid_lines,
-                dotted_lines=self.padding_dashed
+                dotted_lines=self.padding_dashed,
             )
             self.add(self.exterior_rectangle)
             # Add an interior rectangle with no fill color
@@ -67,13 +72,13 @@ class FeatureMap(VGroup):
                 fill_opacity=0.0,
                 width=self.feature_map_size[0] * self.cell_width,
                 height=self.feature_map_size[1] * self.cell_width,
-                stroke_width=self.stroke_width
+                stroke_width=self.stroke_width,
             )
             self.add(self.interior_rectangle)
         else:
             # Just make an exterior rectangle with no dashes.
-            self.untransformed_height = self.feature_map_size[1] * self.cell_width,
+            self.untransformed_height = (self.feature_map_size[1] * self.cell_width,)
-            self.untransformed_width = self.feature_map_size[0] * self.cell_width,
+            self.untransformed_width = (self.feature_map_size[0] * self.cell_width,)
             # Make the exterior rectangle
             self.exterior_rectangle = GriddedRectangle(
                 color=self.color,
@@ -96,6 +101,7 @@ class FeatureMap(VGroup):
         # Sort points through clockwise rotation of a vector in the xy plane
         return self.exterior_rectangle.get_corners_dict()
 
+
 class Convolutional2DLayer(VGroupNeuralNetworkLayer, ThreeDLayer):
     """Handles rendering a convolutional layer for a nn"""
 
@@ -215,7 +221,7 @@ class Convolutional2DLayer(VGroupNeuralNetworkLayer, ThreeDLayer):
                 fill_color=self.color,
                 fill_opacity=self.fill_opacity,
                 padding=self.padding,
-                padding_dashed=self.padding_dashed
+                padding_dashed=self.padding_dashed,
             )
             # Move the feature map
             feature_map.move_to([0, 0, filter_index * self.filter_spacing])
@@ -231,9 +237,7 @@ class Convolutional2DLayer(VGroupNeuralNetworkLayer, ThreeDLayer):
             ApplyMethod(self.feature_maps.set_color, self.color),
         )
 
-    def make_forward_pass_animation(
+    def make_forward_pass_animation(self, run_time=5, layer_args={}, **kwargs):
-        self, run_time=5, layer_args={}, **kwargs
-    ):
         """Convolution forward pass animation"""
         # Note: most of this animation is done in the Convolution3DToConvolution3D layer
         if not self.activation_function is None:

manim_ml/neural_network/layers/convolutional_2d_to_convolutional_2d.py

@@ -3,10 +3,11 @@ import numpy as np
 from manim import *
 from manim_ml.neural_network.layers.convolutional_2d import Convolutional2DLayer
 from manim_ml.neural_network.layers.parent_layers import ConnectiveLayer, ThreeDLayer
-from manim_ml.gridded_rectangle import GriddedRectangle
+from manim_ml.utils.mobjects.gridded_rectangle import GriddedRectangle
 
 from manim.utils.space_ops import rotation_matrix
 
+
 def get_rotated_shift_vectors(input_layer, normalized=False):
     """Rotates the shift vectors"""
     # Make base shift vectors
@@ -24,6 +25,7 @@ def get_rotated_shift_vectors(input_layer, normalized=False):
 
     return right_shift, down_shift
 
+
 class Filters(VGroup):
     """Group for showing a collection of filters connecting two layers"""
 
@@ -144,7 +146,7 @@ class Filters(VGroup):
             )
             # Shift based on the amount of output layer padding
             rectangle.shift(
-                self.output_layer.padding[0] * right_shift, 
+                self.output_layer.padding[0] * right_shift,
             )
             rectangle.shift(
                 self.output_layer.padding[1] * down_shift,
@@ -446,10 +448,7 @@ class Convolutional2DToConvolutional2D(ConnectiveLayer, ThreeDLayer):
             # Do last row move right
             for x_move in range(num_x_moves):
                 # Shift right
-                shift_animation = ApplyMethod(
+                shift_animation = ApplyMethod(filters.shift, self.stride * right_shift)
-                    filters.shift, 
-                    self.stride * right_shift
-                )
                 # shift_animation = self.animate.shift(right_shift)
                 animations.append(shift_animation)
             # Remove the filters
@@ -460,18 +459,14 @@ class Convolutional2DToConvolutional2D(ConnectiveLayer, ThreeDLayer):
                 # Change the output feature map colors
                 change_color_animations = []
                 change_color_animations.append(
-                    ApplyMethod(
+                    ApplyMethod(feature_map.set_color, original_feature_map_color)
-                        feature_map.set_color, 
-                        original_feature_map_color
-                    )
                 )
                 # Change the input feature map colors
                 input_feature_maps = self.input_layer.feature_maps
                 for input_feature_map in input_feature_maps:
                     change_color_animations.append(
                         ApplyMethod(
-                            input_feature_map.set_color, 
+                            input_feature_map.set_color, original_feature_map_color
-                            original_feature_map_color
                         )
                     )
                 # Combine the animations

manim_ml/neural_network/layers/convolutional_2d_to_max_pooling_2d.py

@@ -1,6 +1,6 @@
 import random
 from manim import *
-from manim_ml.gridded_rectangle import GriddedRectangle
+from manim_ml.utils.mobjects.gridded_rectangle import GriddedRectangle
 from manim_ml.neural_network.layers.convolutional_2d_to_convolutional_2d import (
     get_rotated_shift_vectors,
 )
@@ -10,6 +10,7 @@ from manim_ml.neural_network.layers.parent_layers import ConnectiveLayer, ThreeD
 from manim_ml.neural_network.layers.feed_forward import FeedForwardLayer
 from manim_ml.neural_network.layers.convolutional_2d import Convolutional2DLayer
 
+
 class Uncreate(Create):
     def __init__(
         self,
@@ -27,6 +28,7 @@ class Uncreate(Create):
             **kwargs,
         )
 
+
 class Convolutional2DToMaxPooling2D(ConnectiveLayer, ThreeDLayer):
     """Feed Forward to Embedding Layer"""
 
@@ -38,7 +40,7 @@ class Convolutional2DToMaxPooling2D(ConnectiveLayer, ThreeDLayer):
         input_layer: Convolutional2DLayer,
         output_layer: MaxPooling2DLayer,
         active_color=ORANGE,
-        **kwargs
+        **kwargs,
     ):
         super().__init__(input_layer, output_layer, **kwargs)
         self.active_color = active_color
@@ -47,7 +49,7 @@ class Convolutional2DToMaxPooling2D(ConnectiveLayer, ThreeDLayer):
         self,
         input_layer: "NeuralNetworkLayer",
         output_layer: "NeuralNetworkLayer",
-        **kwargs
+        **kwargs,
     ):
         return super().construct_layer(input_layer, output_layer, **kwargs)
 
@@ -94,7 +96,7 @@ class Convolutional2DToMaxPooling2D(ConnectiveLayer, ThreeDLayer):
                         width=cell_width,
                         fill_opacity=0.7,
                         stroke_width=0.0,
-                        z_index=10
+                        z_index=10,
                     )
                     # Move to the correct location
                     kernel_shift_vector = [
@@ -119,10 +121,7 @@ class Convolutional2DToMaxPooling2D(ConnectiveLayer, ThreeDLayer):
                     highlighted_cells.append(cell_rectangle)
             # Rotate the gridded rectangles so they match the angle
             # of the conv maps
-            gridded_rectangle_group = VGroup(
+            gridded_rectangle_group = VGroup(gridded_rectangle, *highlighted_cells)
-                gridded_rectangle, 
-                *highlighted_cells
-            )
             gridded_rectangle_group.rotate(
                 ThreeDLayer.rotation_angle,
                 about_point=gridded_rectangle.get_center(),
@@ -137,26 +136,20 @@ class Convolutional2DToMaxPooling2D(ConnectiveLayer, ThreeDLayer):
             create_rectangle = Create(
                 gridded_rectangle,
             )
-            create_gridded_rectangle_animations.append(
+            create_gridded_rectangle_animations.append(create_rectangle)
-                create_rectangle
-            )
             # 4. Create and fade out highlighted cells
             create_group = AnimationGroup(
                 *[Create(highlighted_cell) for highlighted_cell in highlighted_cells],
-                lag_ratio=1.0
+                lag_ratio=1.0,
             )
             uncreate_group = AnimationGroup(
                 *[Uncreate(highlighted_cell) for highlighted_cell in highlighted_cells],
-                lag_ratio=0.0
+                lag_ratio=0.0,
             )
             create_and_remove_cell_animation = Succession(
-                create_group,
+                create_group, Wait(1.0), uncreate_group
-                Wait(1.0),
-                uncreate_group
-            )
-            create_and_remove_cell_animations.append(
-                create_and_remove_cell_animation
             )
+            create_and_remove_cell_animations.append(create_and_remove_cell_animation)
             # 5. Move and resize the gridded rectangle to the output
             #   feature maps.
             output_gridded_rectangle = GriddedRectangle(
@@ -178,9 +171,10 @@ class Convolutional2DToMaxPooling2D(ConnectiveLayer, ThreeDLayer):
                 self.output_layer.feature_maps[feature_map_index].copy()
             )
             transform_rectangle = ReplacementTransform(
-                gridded_rectangle, output_gridded_rectangle,
+                gridded_rectangle,
+                output_gridded_rectangle,
                 introducer=True,
-                remover=True
+                remover=True,
             )
             transform_gridded_rectangle_animations.append(
                 transform_rectangle,

manim_ml/neural_network/layers/embedding.py

@@ -1,5 +1,5 @@
 from manim import *
-from manim_ml.probability import GaussianDistribution
+from manim_ml.utils.mobjects.probability import GaussianDistribution
 from manim_ml.neural_network.layers.parent_layers import VGroupNeuralNetworkLayer
 
 

manim_ml/neural_network/layers/feed_forward.py

@@ -1,9 +1,12 @@
 from manim import *
 
 from manim_ml.neural_network.activation_functions import get_activation_function_by_name
-from manim_ml.neural_network.activation_functions.activation_function import ActivationFunction
+from manim_ml.neural_network.activation_functions.activation_function import (
+    ActivationFunction,
+)
 from manim_ml.neural_network.layers.parent_layers import VGroupNeuralNetworkLayer
 
+
 class FeedForwardLayer(VGroupNeuralNetworkLayer):
     """Handles rendering a layer for a neural network"""
 

manim_ml/neural_network/layers/image.py

@@ -1,6 +1,6 @@
 from manim import *
 import numpy as np
-from manim_ml.image import GrayscaleImageMobject
+from manim_ml.utils.mobjects.image import GrayscaleImageMobject
 from manim_ml.neural_network.layers.parent_layers import NeuralNetworkLayer
 
 from PIL import Image

manim_ml/neural_network/layers/image_to_convolutional_2d.py

@@ -7,7 +7,7 @@ from manim_ml.neural_network.layers.parent_layers import (
     ThreeDLayer,
     VGroupNeuralNetworkLayer,
 )
-from manim_ml.gridded_rectangle import GriddedRectangle
+from manim_ml.utils.mobjects.gridded_rectangle import GriddedRectangle
 
 
 class ImageToConvolutional2DLayer(VGroupNeuralNetworkLayer, ThreeDLayer):

manim_ml/neural_network/layers/max_pooling_2d.py

@@ -1,5 +1,5 @@
 from manim import *
-from manim_ml.gridded_rectangle import GriddedRectangle
+from manim_ml.utils.mobjects.gridded_rectangle import GriddedRectangle
 
 from manim_ml.neural_network.layers.parent_layers import (
     ThreeDLayer,

manim_ml/neural_network/layers/max_pooling_2d_to_feed_forward.py

@@ -1,8 +1,11 @@
 from manim import *
-from manim_ml.neural_network.layers.convolutional_2d_to_feed_forward import Convolutional2DToFeedForward
+from manim_ml.neural_network.layers.convolutional_2d_to_feed_forward import (
+    Convolutional2DToFeedForward,
+)
 from manim_ml.neural_network.layers.feed_forward import FeedForwardLayer
 from manim_ml.neural_network.layers.max_pooling_2d import MaxPooling2DLayer
 
+
 class MaxPooling2DToFeedForward(Convolutional2DToFeedForward):
     """Feed Forward to Embedding Layer"""
 

manim_ml/neural_network/layers/paired_query.py

@@ -1,6 +1,6 @@
 from manim import *
 from manim_ml.neural_network.layers.parent_layers import NeuralNetworkLayer
-from manim_ml.image import GrayscaleImageMobject, LabeledColorImage
+from manim_ml.utils.mobjects.image import GrayscaleImageMobject, LabeledColorImage
 import numpy as np
 
 

manim_ml/neural_network/layers/triplet.py

@@ -1,6 +1,6 @@
 from manim import *
 from manim_ml.neural_network.layers import NeuralNetworkLayer
-from manim_ml.image import GrayscaleImageMobject, LabeledColorImage
+from manim_ml.utils.mobjects.image import GrayscaleImageMobject, LabeledColorImage
 import numpy as np
 
 

manim_ml/neural_network/layers/util.py

@@ -4,6 +4,7 @@ from manim import *
 from manim_ml.neural_network.layers.parent_layers import BlankConnective, ThreeDLayer
 from manim_ml.neural_network.layers import connective_layers_list
 
+
 def get_connective_layer(input_layer, output_layer):
     """
     Deduces the relevant connective layer
@@ -12,8 +13,9 @@ def get_connective_layer(input_layer, output_layer):
     for candidate_class in connective_layers_list:
         input_class = candidate_class.input_class
         output_class = candidate_class.output_class
-        if isinstance(input_layer, input_class) and \
+        if isinstance(input_layer, input_class) and isinstance(
-            isinstance(output_layer, output_class):
+            output_layer, output_class
+        ):
             connective_layer_class = candidate_class
             break
 

manim_ml/neural_network/neural_network.py

@@ -17,12 +17,13 @@ from manim import *
 
 from manim_ml.neural_network.layers.parent_layers import ConnectiveLayer, ThreeDLayer
 from manim_ml.neural_network.layers.util import get_connective_layer
-from manim_ml.list_group import ListGroup
+from manim_ml.utils.mobjects.list_group import ListGroup
-from manim_ml.neural_network.neural_network_transformations import (
+from manim_ml.neural_network.animations.neural_network_transformations import (
     InsertLayer,
     RemoveLayer,
 )
 
+
 class NeuralNetwork(Group):
     """Neural Network Visualization Container Class"""
 
@@ -59,10 +60,7 @@ class NeuralNetwork(Group):
         # Make the connective layers
         self.connective_layers, self.all_layers = self._construct_connective_layers()
         # Make overhead title
-        self.title = Text(
+        self.title = Text(self.title_text, font_size=DEFAULT_FONT_SIZE / 2)
-            self.title_text, 
-            font_size=DEFAULT_FONT_SIZE / 2
-        )
         self.title.next_to(self, UP, 1.0)
         self.add(self.title)
         # Place layers at correct z index
@@ -92,11 +90,11 @@ class NeuralNetwork(Group):
             raise Exception(f"Uncrecognized input layers type: {type(input_layers)}")
 
     def add_connection(
-        self, 
+        self,
-        start_layer_name, 
+        start_layer_name,
-        end_layer_name, 
+        end_layer_name,
         connection_style="default",
-        connection_position="bottom"
+        connection_position="bottom",
     ):
         """Add connection from start layer to end layer"""
         assert connection_style in ["default"]
@@ -106,7 +104,7 @@ class NeuralNetwork(Group):
             connection = NetworkConnection(
                 self.input_layers_dict[start_layer_name],
                 self.input_layers_dict[end_layer_name],
-                arc_direction="down" # TODO generalize this more
+                arc_direction="down",  # TODO generalize this more
             )
             self.connections.append(connection)
             self.add(connection)
@@ -138,20 +136,20 @@ class NeuralNetwork(Group):
             current_layer.move_to(previous_layer.get_center())
 
             if layout_direction == "left_to_right":
-                x_shift = previous_layer.get_width() / 2 \
+                x_shift = (
-                        + current_layer.get_width() / 2 \
+                    previous_layer.get_width() / 2
-                        + self.layer_spacing
+                    + current_layer.get_width() / 2
+                    + self.layer_spacing
+                )
                 shift_vector = np.array([x_shift, 0, 0])
             elif layout_direction == "top_to_bottom":
-                y_shift = -((
+                y_shift = -(
-                    previous_layer.get_width() / 2 \
+                    (previous_layer.get_width() / 2 + current_layer.get_width() / 2)
-                    + current_layer.get_width() / 2
+                    + self.layer_spacing
-                ) + self.layer_spacing)
+                )
                 shift_vector = np.array([0, y_shift, 0])
             else:
-                raise Exception(
+                raise Exception(f"Unrecognized layout direction: {layout_direction}")
-                    f"Unrecognized layout direction: {layout_direction}"
-                )
             current_layer.shift(shift_vector)
 
         # After all layers have been placed place their activation functions
@@ -159,13 +157,15 @@ class NeuralNetwork(Group):
             # Place activation function
             if hasattr(current_layer, "activation_function"):
                 if not current_layer.activation_function is None:
-                    up_movement = np.array([
+                    up_movement = np.array(
-                        0,
+                        [
-                        current_layer.get_height() / 2
+                            0,
-                        + current_layer.activation_function.get_height() / 2
+                            current_layer.get_height() / 2
-                        + 0.5 * self.layer_spacing,
+                            + current_layer.activation_function.get_height() / 2
-                        0,
+                            + 0.5 * self.layer_spacing,
-                    ])
+                            0,
+                        ]
+                    )
                     current_layer.activation_function.move_to(
                         current_layer,
                     )
@@ -259,9 +259,7 @@ class NeuralNetwork(Group):
                     current_layer_args = layer_args[layer]
             # Perform the forward pass of the current layer
             layer_forward_pass = layer.make_forward_pass_animation(
-                layer_args=current_layer_args, 
+                layer_args=current_layer_args, run_time=per_layer_runtime, **kwargs
-                run_time=per_layer_runtime, 
-                **kwargs
             )
             # Animate a forward pass for incoming connections
             connection_input_pass = AnimationGroup()
@@ -272,14 +270,12 @@ class NeuralNetwork(Group):
                         connection_input_pass = ShowPassingFlash(
                             connection,
                             run_time=layer_forward_pass.run_time,
-                            time_width=0.2
+                            time_width=0.2,
                         )
                         break
 
             layer_forward_pass = AnimationGroup(
-                layer_forward_pass,
+                layer_forward_pass, connection_input_pass, lag_ratio=0.0
-                connection_input_pass,
-                lag_ratio=0.0
             )
             all_animations.append(layer_forward_pass)
         # Make the animation group

manim_ml/one_to_one_sync.py

@@ -1,11 +0,0 @@
-"""
-    Module for handling syncing two animations one to one.
-    The goal here is to zip up two classes and their respective animations,
-    and create a joint class with the same animations that runs the animations
-    for both classes at the same time. This way we can connect two isomorphic
-    views of the same concept and visualize them at the same time. 
-"""
-
-
-class OneToOneSync:
-    pass

manim_ml/utils/mobjects/connections.py

@@ -1,30 +1,27 @@
 import numpy as np
 from manim import *
 
+
 class NetworkConnection(VGroup):
     """
-        This class allows for creating connections
+    This class allows for creating connections
-        between locations in a network
+    between locations in a network
     """
-    direction_vector_map = {
+
-        "up": UP,
+    direction_vector_map = {"up": UP, "down": DOWN, "left": LEFT, "right": RIGHT}
-        "down": DOWN,
-        "left": LEFT,
-        "right": RIGHT
-    }
 
     def __init__(
-        self, 
+        self,
-        start_mobject, 
+        start_mobject,
-        end_mobject, 
+        end_mobject,
         arc_direction="straight",
         buffer=0.05,
         arc_distance=0.3,
         stroke_width=2.0,
         color=WHITE,
-        active_color=ORANGE
+        active_color=ORANGE,
     ):
-        """Creates an arrow with right angles in it connecting 
+        """Creates an arrow with right angles in it connecting
         two mobjects.
 
         Parameters
@@ -72,34 +69,28 @@ class NetworkConnection(VGroup):
             # Make an arrow
             arrow_line = Line(
                 left_mobject.get_right() + np.array([self.buffer, 0.0, 0.0]),
-                right_mobject.get_left() + np.array([-1 * self.buffer, 0.0, 0.0])
+                right_mobject.get_left() + np.array([-1 * self.buffer, 0.0, 0.0]),
-            )
-            arrow = Arrow(
-                arrow_line,
-                color=self.color,
-                stroke_width=self.stroke_width
             )
+            arrow = Arrow(arrow_line, color=self.color, stroke_width=self.stroke_width)
             self.straight_arrow = arrow
             self.add(arrow)
         else:
             # Figure out the direction of the arc
-            direction_vector = NetworkConnection.direction_vector_map[self.arc_direction]
+            direction_vector = NetworkConnection.direction_vector_map[
+                self.arc_direction
+            ]
             # Make the start arc piece
-            start_line_start = left_mobject.get_critical_point(
+            start_line_start = left_mobject.get_critical_point(direction_vector)
-                direction_vector
-            )
             start_line_start += direction_vector * self.buffer
             start_line_end = start_line_start + direction_vector * self.arc_distance
             self.start_line = Line(
                 start_line_start,
                 start_line_end,
                 color=self.color,
-                stroke_width=self.stroke_width
+                stroke_width=self.stroke_width,
             )
             # Make the end arc piece with an arrow
-            end_line_end = right_mobject.get_critical_point(
+            end_line_end = right_mobject.get_critical_point(direction_vector)
-                direction_vector
-            )
             end_line_end += direction_vector * self.buffer
             end_line_start = end_line_end + direction_vector * self.arc_distance
             self.end_arrow = Arrow(
@@ -108,14 +99,14 @@ class NetworkConnection(VGroup):
                 color=WHITE,
                 fill_color=WHITE,
                 stroke_opacity=1.0,
-                buff=0.0
+                buff=0.0,
             )
             # Make the middle arc piece
             self.middle_line = Line(
                 start_line_end,
                 end_line_start,
                 color=self.color,
-                stroke_width=self.stroke_width
+                stroke_width=self.stroke_width,
             )
             # Add the mobjects
             self.add(
@@ -130,23 +121,23 @@ class NetworkConnection(VGroup):
         if self.arc_direction == "straight":
             return ShowPassingFlash(
                 self.straight_arrow.copy().set_color(self.active_color),
-                time_width=time_width
+                time_width=time_width,
             )
         else:
             # Animate the start line
             start_line_animation = ShowPassingFlash(
                 self.start_line.copy().set_color(self.active_color),
-                time_width=time_width
+                time_width=time_width,
             )
             # Animate the middle line
             middle_line_animation = ShowPassingFlash(
                 self.middle_line.copy().set_color(self.active_color),
-                time_width=time_width
+                time_width=time_width,
             )
             # Animate the end line
             end_line_animation = ShowPassingFlash(
                 self.end_arrow.copy().set_color(self.active_color),
-                time_width=time_width
+                time_width=time_width,
             )
 
             return AnimationGroup(
@@ -154,5 +145,5 @@ class NetworkConnection(VGroup):
                 middle_line_animation,
                 end_line_animation,
                 lag_ratio=1.0,
-                run_time=run_time
+                run_time=run_time,
-            )
+            )

manim_ml/utils/mobjects/gridded_rectangle.py

@@ -14,7 +14,7 @@ class GriddedRectangle(VGroup):
         close_new_points=True,
         grid_xstep=None,
         grid_ystep=None,
-        grid_stroke_width=0.0, # DEFAULT_STROKE_WIDTH/2,
+        grid_stroke_width=0.0,  # DEFAULT_STROKE_WIDTH/2,
         grid_stroke_color=ORANGE,
         grid_stroke_opacity=1.0,
         stroke_width=2.0,
@@ -48,7 +48,7 @@ class GriddedRectangle(VGroup):
                 fill_color=color,
                 stroke_opacity=0.0,
                 fill_opacity=fill_opacity,
-                shade_in_3d=True
+                shade_in_3d=True,
             )
             self.rectangle = no_border_rectangle
             border_rectangle = Rectangle(
@@ -58,7 +58,7 @@ class GriddedRectangle(VGroup):
                 fill_color=color,
                 fill_opacity=fill_opacity,
                 shade_in_3d=True,
-                stroke_width=stroke_width
+                stroke_width=stroke_width,
             )
             self.dotted_lines = DashedVMobject(
                 border_rectangle,
@@ -73,7 +73,7 @@ class GriddedRectangle(VGroup):
                 stroke_width=stroke_width,
                 fill_color=color,
                 fill_opacity=fill_opacity,
-                shade_in_3d=True
+                shade_in_3d=True,
             )
         self.add(self.rectangle)
         # Make grid lines
@@ -123,7 +123,7 @@ class GriddedRectangle(VGroup):
                         stroke_color=self.grid_stroke_color,
                         stroke_width=self.grid_stroke_width,
                         stroke_opacity=self.grid_stroke_opacity,
-                        shade_in_3d=True
+                        shade_in_3d=True,
                     )
                     for i in range(1, count)
                 )

manim_ml/utils/testing/frames_comparison.py

@@ -27,6 +27,7 @@ _tests_root_dir_path = Path(__file__).absolute().parents[2]
 print(f"Tests root path: {_tests_root_dir_path}")
 PATH_CONTROL_DATA = _tests_root_dir_path / Path("control_data", "graphical_units_data")
 
+
 def frames_comparison(
     func=None,
     *,
@@ -254,4 +255,4 @@ def _config_test(last_frame: bool) -> ManimConfig:
                 else "config_graphical_tests_multiframes.cfg"
             ),
         ),
-    )
+    )

setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="manim_ml",
-    version="0.0.15",
+    version="0.0.16",
     description=(" Machine Learning Animations in python using Manim."),
     packages=find_packages(),
 )

tests/conftest.py

@@ -1,5 +1,3 @@
-
-
 from __future__ import annotations
 
 import os
@@ -9,6 +7,7 @@ import pytest
 
 from manim import config, tempconfig
 
+
 def pytest_addoption(parser):
     parser.addoption(
         "--skip_slow",

tests/test_activation_function.py

@@ -13,15 +13,19 @@ config.pixel_width = 1900
 config.frame_height = 6.0
 config.frame_width = 6.0
 
+
 class CombinedScene(ThreeDScene):
     def construct(self):
         image = Image.open("../assets/mnist/digit.jpeg")
         numpy_image = np.asarray(image)
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 ImageLayer(numpy_image, height=1.5),
                 Convolutional2DLayer(1, 7, filter_spacing=0.32),
-                Convolutional2DLayer(3, 5, 3, filter_spacing=0.32, activation_function="ReLU"),
+                Convolutional2DLayer(
+                    3, 5, 3, filter_spacing=0.32, activation_function="ReLU"
+                ),
                 FeedForwardLayer(3, activation_function="Sigmoid"),
             ],
             layer_spacing=0.25,
@@ -32,4 +36,4 @@ class CombinedScene(ThreeDScene):
         # Play animation
         forward_pass = nn.make_forward_pass_animation()
         self.wait(1)
-        self.play(forward_pass, run_time=30)
+        self.play(forward_pass, run_time=30)

tests/test_camera_move.py

@@ -12,14 +12,11 @@ config.pixel_width = 1900
 config.frame_height = 6.0
 config.frame_width = 6.0
 
+
 class NeuralNetworkScene(ThreeDScene):
     """Test Scene for the Neural Network"""
 
-    def play_camera_follow_forward_pass(
+    def play_camera_follow_forward_pass(self, neural_network, buffer=0.1):
-        self, 
-        neural_network,
-        buffer=0.1
-    ):
         per_layer_animations = neural_network.make_forward_pass_animation(
             return_per_layer_animations=True
         )
@@ -32,7 +29,7 @@ class NeuralNetworkScene(ThreeDScene):
             current_layer = all_layers[layer_index]
             next_layer = all_layers[layer_index + 1]
             group = Group(prev_layer, current_layer, next_layer)
-            
+
             max_width = max(max_width, group.width)
             max_height = max(max_height, group.height)
 
@@ -46,7 +43,8 @@ class NeuralNetworkScene(ThreeDScene):
         # Make the Layer object
         image = Image.open("../assets/mnist/digit.jpeg")
         numpy_image = np.asarray(image)
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 ImageLayer(numpy_image, height=1.5),
                 Convolutional2DLayer(1, 7, filter_spacing=0.32),
                 Convolutional2DLayer(3, 5, 3, filter_spacing=0.32),
@@ -60,4 +58,4 @@ class NeuralNetworkScene(ThreeDScene):
         # Make Animation
         self.add(nn)
         # self.play(Create(nn))
-        self.play_camera_follow_forward_pass(nn)
+        self.play_camera_follow_forward_pass(nn)

tests/test_conv_padding.py

@@ -14,23 +14,24 @@ config.pixel_width = 1900
 config.frame_height = 6.0
 config.frame_width = 6.0
 
-class CombinedScene(ThreeDScene):
 
+class CombinedScene(ThreeDScene):
     def construct(self):
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 Convolutional2DLayer(
-                    num_feature_maps=1, 
+                    num_feature_maps=1,
-                    feature_map_size=7, 
+                    feature_map_size=7,
                     padding=1,
-                    padding_dashed=True
+                    padding_dashed=True,
                 ),
                 Convolutional2DLayer(
-                    num_feature_maps=3, 
+                    num_feature_maps=3,
-                    feature_map_size=7, 
+                    feature_map_size=7,
                     filter_size=3,
                     padding=0,
-                    padding_dashed=False
+                    padding_dashed=False,
                 ),
                 FeedForwardLayer(3),
             ],
@@ -44,23 +45,22 @@ class CombinedScene(ThreeDScene):
         self.wait(1)
         self.play(forward_pass, run_time=30)
 
+
 @frames_comparison
 def test_ConvPadding(scene):
     # Make nn
-    nn = NeuralNetwork([
+    nn = NeuralNetwork(
+        [
             Convolutional2DLayer(
-                num_feature_maps=1, 
+                num_feature_maps=1, feature_map_size=7, padding=1, padding_dashed=True
-                feature_map_size=7, 
-                padding=1,
-                padding_dashed=True
             ),
             Convolutional2DLayer(
-                num_feature_maps=3, 
+                num_feature_maps=3,
-                feature_map_size=7, 
+                feature_map_size=7,
                 filter_size=3,
                 padding=1,
                 filter_spacing=0.35,
-                padding_dashed=False
+                padding_dashed=False,
             ),
             FeedForwardLayer(3),
         ],
@@ -71,4 +71,4 @@ def test_ConvPadding(scene):
     scene.add(nn)
     # Play animation
     forward_pass = nn.make_forward_pass_animation()
-    scene.play(forward_pass, run_time=30)
+    scene.play(forward_pass, run_time=30)

tests/test_convolutional_2d_layer.py

@@ -6,6 +6,7 @@ from manim_ml.neural_network.layers.feed_forward import FeedForwardLayer
 from manim_ml.neural_network.layers.image import ImageLayer
 from manim_ml.neural_network.neural_network import NeuralNetwork
 
+
 class SingleConvolutionalLayerScene(ThreeDScene):
     def construct(self):
         # Make nn
@@ -21,6 +22,7 @@ class SingleConvolutionalLayerScene(ThreeDScene):
         )
         # self.play(nn.make_forward_pass_animation(run_time=5))
 
+
 class Simple3DConvScene(ThreeDScene):
     def construct(self):
         """
@@ -33,16 +35,8 @@ class Simple3DConvScene(ThreeDScene):
         """
         # Make nn
         layers = [
-            Convolutional2DLayer(
+            Convolutional2DLayer(num_feature_maps=1, feature_map_size=3, filter_size=3),
-                num_feature_maps=1,
+            Convolutional2DLayer(num_feature_maps=1, feature_map_size=3, filter_size=3),
-                feature_map_size=3,
-                filter_size=3
-            ),
-            Convolutional2DLayer(
-                num_feature_maps=1,
-                feature_map_size=3,
-                filter_size=3
-            ),
         ]
         nn = NeuralNetwork(layers)
         # Center the nn
@@ -52,18 +46,21 @@ class Simple3DConvScene(ThreeDScene):
         # self.set_camera_orientation(phi=280*DEGREES, theta=-10*DEGREES, gamma=90*DEGREES)
         self.play(nn.make_forward_pass_animation(run_time=30))
 
+
 # Make the specific scene
 config.pixel_height = 1200
 config.pixel_width = 1900
 config.frame_height = 6.0
 config.frame_width = 6.0
 
+
 class CombinedScene(ThreeDScene):
     def construct(self):
         image = Image.open("../assets/mnist/digit.jpeg")
         numpy_image = np.asarray(image)
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 ImageLayer(numpy_image, height=1.5),
                 Convolutional2DLayer(1, 7, filter_spacing=0.32),
                 Convolutional2DLayer(3, 5, 3, filter_spacing=0.32),
@@ -79,4 +76,4 @@ class CombinedScene(ThreeDScene):
         # Play animation
         forward_pass = nn.make_forward_pass_animation()
         self.wait(1)
-        self.play(forward_pass)
+        self.play(forward_pass)

tests/test_decision_tree.py

@@ -9,6 +9,7 @@ from sklearn import datasets
 import sklearn
 import matplotlib.pyplot as plt
 
+
 def learn_iris_decision_tree(iris):
     decision_tree = DecisionTreeClassifier(
         random_state=1, max_depth=3, max_leaf_nodes=6
@@ -17,11 +18,13 @@ def learn_iris_decision_tree(iris):
     # output the decisioin tree in some format
     return decision_tree
 
+
 def make_sklearn_tree(dataset, max_tree_depth=3):
     tree = learn_iris_decision_tree(dataset)
     feature_names = dataset.feature_names[0:2]
     return tree, tree.tree_
 
+
 class DecisionTreeScene(Scene):
     def construct(self):
         """Makes a decision tree object"""
@@ -44,6 +47,7 @@ class DecisionTreeScene(Scene):
         self.play(create_decision_tree)
         # self.play(create_decision_tree)
 
+
 class SurfacePlot(Scene):
     def construct(self):
         iris_dataset = datasets.load_iris()

tests/test_feed_forward.py

@@ -5,26 +5,21 @@ from manim_ml.neural_network import NeuralNetwork, FeedForwardLayer
 
 __module_test__ = "feed_forward"
 
+
 @frames_comparison
 def test_FeedForwardScene(scene):
     """Tests the appearance of a feed forward network"""
-    nn = NeuralNetwork([
+    nn = NeuralNetwork([FeedForwardLayer(3), FeedForwardLayer(5), FeedForwardLayer(3)])
-        FeedForwardLayer(3),
-        FeedForwardLayer(5),
-        FeedForwardLayer(3)
-    ])
 
     scene.add(nn)
 
-class FeedForwardScene(Scene):
 
+class FeedForwardScene(Scene):
     def construct(self):
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
-            FeedForwardLayer(3),
+            [FeedForwardLayer(3), FeedForwardLayer(5), FeedForwardLayer(3)]
-            FeedForwardLayer(5),
+        )
-            FeedForwardLayer(3)
-        ])
 
         self.add(nn)
 
-        self.play(nn.make_forward_pass_animation())
+        self.play(nn.make_forward_pass_animation())

tests/test_gridded_rectangle.py

@@ -1,12 +1,8 @@
 from manim import *
-from manim_ml.gridded_rectangle import GriddedRectangle
+from manim_ml.utils.mobjects.gridded_rectangle import GriddedRectangle
+
 
 class TestGriddedRectangleScene(ThreeDScene):
-
     def construct(self):
-        rect = GriddedRectangle(
+        rect = GriddedRectangle(color=ORANGE, width=3, height=3)
-            color=ORANGE,
+        self.add(rect)
-            width=3,
-            height=3
-        )
-        self.add(rect)

tests/test_image_homotopy.py

@@ -1,7 +1,7 @@
 from PIL import Image
 
 from manim import *
-from manim_ml.image import GrayscaleImageMobject
+from manim_ml.utils.mobjects.image import GrayscaleImageMobject
 from manim_ml.neural_network.layers.parent_layers import ThreeDLayer
 
 

tests/test_max_pool.py

@@ -14,12 +14,14 @@ config.pixel_width = 1900
 config.frame_height = 6.0
 config.frame_width = 6.0
 
+
 class CombinedScene(ThreeDScene):
     def construct(self):
         image = Image.open("../assets/mnist/digit.jpeg")
         numpy_image = np.asarray(image)
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 ImageLayer(numpy_image, height=1.5),
                 Convolutional2DLayer(1, 8, filter_spacing=0.32),
                 Convolutional2DLayer(3, 6, 3, filter_spacing=0.32),
@@ -37,12 +39,14 @@ class CombinedScene(ThreeDScene):
         self.wait(1)
         self.play(forward_pass)
 
+
 class SmallNetwork(ThreeDScene):
     def construct(self):
         image = Image.open("../assets/mnist/digit.jpeg")
         numpy_image = np.asarray(image)
         # Make nn
-        nn = NeuralNetwork([
+        nn = NeuralNetwork(
+            [
                 ImageLayer(numpy_image, height=1.5),
                 Convolutional2DLayer(1, 8, filter_spacing=0.32),
                 MaxPooling2DLayer(kernel_size=2),
@@ -55,4 +59,4 @@ class SmallNetwork(ThreeDScene):
         # Play animation
         forward_pass = nn.make_forward_pass_animation()
         self.wait(1)
-        self.play(forward_pass)
+        self.play(forward_pass)

tests/test_mcmc.py

@@ -1,31 +1,33 @@
 from manim import *
-from manim_ml.diffusion.mcmc import MCMCAxes, MultidimensionalGaussianPosterior, metropolis_hastings_sampler
+from manim_ml.diffusion.mcmc import (
+    MCMCAxes,
+    MultidimensionalGaussianPosterior,
+    metropolis_hastings_sampler,
+)
+
 # Make the specific scene
 config.pixel_height = 1200
 config.pixel_width = 1200
 config.frame_height = 12.0
 config.frame_width = 12.0
 
+
 def test_metropolis_hastings_sampler(iterations=100):
     samples, _, candidates = metropolis_hastings_sampler(iterations=iterations)
     assert samples.shape == (iterations, 2)
 
-class MCMCTest(Scene):
 
+class MCMCTest(Scene):
     def construct(self):
         axes = MCMCAxes()
         self.play(Create(axes))
         gaussian_posterior = MultidimensionalGaussianPosterior(
-            mu=np.array([0.0, 0.0]),
+            mu=np.array([0.0, 0.0]), var=np.array([4.0, 2.0])
-            var=np.array([4.0, 2.0])
-        )
-        show_gaussian_animation = axes.show_ground_truth_gaussian(
-            gaussian_posterior
         )
+        show_gaussian_animation = axes.show_ground_truth_gaussian(gaussian_posterior)
         self.play(show_gaussian_animation)
         chain_sampling_animation = axes.visualize_metropolis_hastings_chain_sampling(
-            log_prob_fn=gaussian_posterior,
+            log_prob_fn=gaussian_posterior, sampling_kwargs={"iterations": 1000}
-            sampling_kwargs={"iterations": 1000}
         )
 
         self.play(chain_sampling_animation)

tests/test_nested_neural_networks.py

@@ -1,4 +1,4 @@
 """
     The purpose of this test is to ensure that it is possible
     to have nested neural network layers. 
-"""
+"""

tests/test_residual_connection.py

@@ -9,49 +9,53 @@ import numpy as np
 
 __module_test__ = "residual"
 
+
 @frames_comparison
 def test_ResidualConnectionScene(scene):
     """Tests the appearance of a residual connection"""
-    nn = NeuralNetwork({
+    nn = NeuralNetwork(
-        "layer1": FeedForwardLayer(3),
+        {
-        "layer2": FeedForwardLayer(5),
+            "layer1": FeedForwardLayer(3),
-        "layer3": FeedForwardLayer(3)
+            "layer2": FeedForwardLayer(5),
-    })
+            "layer3": FeedForwardLayer(3),
+        }
+    )
 
     scene.add(nn)
 
+
 # Make the specific scene
 config.pixel_height = 1200
 config.pixel_width = 1900
 config.frame_height = 6.0
 config.frame_width = 6.0
 
-class FeedForwardScene(Scene):
 
+class FeedForwardScene(Scene):
     def construct(self):
-        nn = NeuralNetwork({
+        nn = NeuralNetwork(
-            "layer1": FeedForwardLayer(4),
+            {
-            "layer2": FeedForwardLayer(4),
+                "layer1": FeedForwardLayer(4),
-            "layer3": FeedForwardLayer(4)
+                "layer2": FeedForwardLayer(4),
-        },
+                "layer3": FeedForwardLayer(4),
-        layer_spacing=0.45)
+            },
+            layer_spacing=0.45,
+        )
 
         nn.add_connection("layer1", "layer3")
 
         self.add(nn)
 
-        self.play(
+        self.play(nn.make_forward_pass_animation(), run_time=8)
-            nn.make_forward_pass_animation(),
+
-            run_time=8
-        )
 
 class ConvScene(ThreeDScene):
-
     def construct(self):
         image = Image.open("../assets/mnist/digit.jpeg")
         numpy_image = np.asarray(image)
 
-        nn = NeuralNetwork({
+        nn = NeuralNetwork(
+            {
                 "layer1": Convolutional2DLayer(1, 5, padding=1),
                 "layer2": Convolutional2DLayer(1, 5, 3, padding=1),
                 "layer3": Convolutional2DLayer(1, 5, 3, padding=1),
@@ -63,7 +67,4 @@ class ConvScene(ThreeDScene):
 
         self.add(nn)
 
-        self.play(
+        self.play(nn.make_forward_pass_animation(), run_time=8)
-            nn.make_forward_pass_animation(),
-            run_time=8
-        )

tests/test_show_gaussian.py

@@ -1,14 +1,12 @@
 from manim import *
-from manim_ml.probability import GaussianDistribution
+from manim_ml.utils.mobjects.probability import GaussianDistribution
+
 
 class TestShowGaussian(Scene):
-
     def construct(self):
         axes = Axes()
         self.add(axes)
         gaussian = GaussianDistribution(
-            axes,
+            axes, mean=np.array([0.0, 0.0]), cov=np.array([[2.0, 0.0], [0.0, 1.0]])
-            mean=np.array([0.0, 0.0]),
-            cov=np.array([[2.0, 0.0], [0.0, 1.0]])
         )
         self.add(gaussian)

tests/test_succession.py

@@ -1,7 +1,7 @@
 from manim import *
 
-class TestSuccession(Scene):
 
+class TestSuccession(Scene):
     def construct(self):
         white_dot = Dot(color=WHITE)
         white_dot.shift(UP)
@@ -16,4 +16,4 @@ class TestSuccession(Scene):
                 Wait(1),
                 Uncreate(red_dot),
             )
-        )
+        )