Add files using upload-large-folder tool

pythonProject/.venv/Lib/site-packages/networkx/linalg/tests/test_graphmatrix.py

@@ -0,0 +1,276 @@
+import pytest
+
+np = pytest.importorskip("numpy")
+pytest.importorskip("scipy")
+
+import networkx as nx
+from networkx.exception import NetworkXError
+from networkx.generators.degree_seq import havel_hakimi_graph
+
+
+def test_incidence_matrix_simple():
+    deg = [3, 2, 2, 1, 0]
+    G = havel_hakimi_graph(deg)
+    deg = [(1, 0), (1, 0), (1, 0), (2, 0), (1, 0), (2, 1), (0, 1), (0, 1)]
+    MG = nx.random_clustered_graph(deg, seed=42)
+
+    I = nx.incidence_matrix(G, dtype=int).todense()
+    # fmt: off
+    expected = np.array(
+        [[1, 1, 1, 0],
+         [0, 1, 0, 1],
+         [1, 0, 0, 1],
+         [0, 0, 1, 0],
+         [0, 0, 0, 0]]
+    )
+    # fmt: on
+    np.testing.assert_equal(I, expected)
+
+    I = nx.incidence_matrix(MG, dtype=int).todense()
+    # fmt: off
+    expected = np.array(
+        [[1, 0, 0, 0, 0, 0, 0],
+         [1, 0, 0, 0, 0, 0, 0],
+         [0, 1, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 0],
+         [0, 1, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 1, 1, 0],
+         [0, 0, 0, 0, 0, 1, 1],
+         [0, 0, 0, 0, 1, 0, 1]]
+    )
+    # fmt: on
+    np.testing.assert_equal(I, expected)
+
+    with pytest.raises(NetworkXError):
+        nx.incidence_matrix(G, nodelist=[0, 1])
+
+
+class TestGraphMatrix:
+    @classmethod
+    def setup_class(cls):
+        deg = [3, 2, 2, 1, 0]
+        cls.G = havel_hakimi_graph(deg)
+        # fmt: off
+        cls.OI = np.array(
+            [[-1, -1, -1, 0],
+             [1, 0, 0, -1],
+             [0, 1, 0, 1],
+             [0, 0, 1, 0],
+             [0, 0, 0, 0]]
+        )
+        cls.A = np.array(
+            [[0, 1, 1, 1, 0],
+             [1, 0, 1, 0, 0],
+             [1, 1, 0, 0, 0],
+             [1, 0, 0, 0, 0],
+             [0, 0, 0, 0, 0]]
+        )
+        # fmt: on
+        cls.WG = havel_hakimi_graph(deg)
+        cls.WG.add_edges_from(
+            (u, v, {"weight": 0.5, "other": 0.3}) for (u, v) in cls.G.edges()
+        )
+        # fmt: off
+        cls.WA = np.array(
+            [[0, 0.5, 0.5, 0.5, 0],
+             [0.5, 0, 0.5, 0, 0],
+             [0.5, 0.5, 0, 0, 0],
+             [0.5, 0, 0, 0, 0],
+             [0, 0, 0, 0, 0]]
+        )
+        # fmt: on
+        cls.MG = nx.MultiGraph(cls.G)
+        cls.MG2 = cls.MG.copy()
+        cls.MG2.add_edge(0, 1)
+        # fmt: off
+        cls.MG2A = np.array(
+            [[0, 2, 1, 1, 0],
+             [2, 0, 1, 0, 0],
+             [1, 1, 0, 0, 0],
+             [1, 0, 0, 0, 0],
+             [0, 0, 0, 0, 0]]
+        )
+        cls.MGOI = np.array(
+            [[-1, -1, -1, -1, 0],
+             [1, 1, 0, 0, -1],
+             [0, 0, 1, 0, 1],
+             [0, 0, 0, 1, 0],
+             [0, 0, 0, 0, 0]]
+        )
+        # fmt: on
+        cls.no_edges_G = nx.Graph([(1, 2), (3, 2, {"weight": 8})])
+        cls.no_edges_A = np.array([[0, 0], [0, 0]])
+
+    def test_incidence_matrix(self):
+        "Conversion to incidence matrix"
+        I = nx.incidence_matrix(
+            self.G,
+            nodelist=sorted(self.G),
+            edgelist=sorted(self.G.edges()),
+            oriented=True,
+            dtype=int,
+        ).todense()
+        np.testing.assert_equal(I, self.OI)
+
+        I = nx.incidence_matrix(
+            self.G,
+            nodelist=sorted(self.G),
+            edgelist=sorted(self.G.edges()),
+            oriented=False,
+            dtype=int,
+        ).todense()
+        np.testing.assert_equal(I, np.abs(self.OI))
+
+        I = nx.incidence_matrix(
+            self.MG,
+            nodelist=sorted(self.MG),
+            edgelist=sorted(self.MG.edges()),
+            oriented=True,
+            dtype=int,
+        ).todense()
+        np.testing.assert_equal(I, self.OI)
+
+        I = nx.incidence_matrix(
+            self.MG,
+            nodelist=sorted(self.MG),
+            edgelist=sorted(self.MG.edges()),
+            oriented=False,
+            dtype=int,
+        ).todense()
+        np.testing.assert_equal(I, np.abs(self.OI))
+
+        I = nx.incidence_matrix(
+            self.MG2,
+            nodelist=sorted(self.MG2),
+            edgelist=sorted(self.MG2.edges()),
+            oriented=True,
+            dtype=int,
+        ).todense()
+        np.testing.assert_equal(I, self.MGOI)
+
+        I = nx.incidence_matrix(
+            self.MG2,
+            nodelist=sorted(self.MG),
+            edgelist=sorted(self.MG2.edges()),
+            oriented=False,
+            dtype=int,
+        ).todense()
+        np.testing.assert_equal(I, np.abs(self.MGOI))
+
+        I = nx.incidence_matrix(self.G, dtype=np.uint8)
+        assert I.dtype == np.uint8
+
+    def test_weighted_incidence_matrix(self):
+        I = nx.incidence_matrix(
+            self.WG,
+            nodelist=sorted(self.WG),
+            edgelist=sorted(self.WG.edges()),
+            oriented=True,
+            dtype=int,
+        ).todense()
+        np.testing.assert_equal(I, self.OI)
+
+        I = nx.incidence_matrix(
+            self.WG,
+            nodelist=sorted(self.WG),
+            edgelist=sorted(self.WG.edges()),
+            oriented=False,
+            dtype=int,
+        ).todense()
+        np.testing.assert_equal(I, np.abs(self.OI))
+
+        # np.testing.assert_equal(nx.incidence_matrix(self.WG,oriented=True,
+        #                                  weight='weight').todense(),0.5*self.OI)
+        # np.testing.assert_equal(nx.incidence_matrix(self.WG,weight='weight').todense(),
+        #              np.abs(0.5*self.OI))
+        # np.testing.assert_equal(nx.incidence_matrix(self.WG,oriented=True,weight='other').todense(),
+        #              0.3*self.OI)
+
+        I = nx.incidence_matrix(
+            self.WG,
+            nodelist=sorted(self.WG),
+            edgelist=sorted(self.WG.edges()),
+            oriented=True,
+            weight="weight",
+        ).todense()
+        np.testing.assert_equal(I, 0.5 * self.OI)
+
+        I = nx.incidence_matrix(
+            self.WG,
+            nodelist=sorted(self.WG),
+            edgelist=sorted(self.WG.edges()),
+            oriented=False,
+            weight="weight",
+        ).todense()
+        np.testing.assert_equal(I, np.abs(0.5 * self.OI))
+
+        I = nx.incidence_matrix(
+            self.WG,
+            nodelist=sorted(self.WG),
+            edgelist=sorted(self.WG.edges()),
+            oriented=True,
+            weight="other",
+        ).todense()
+        np.testing.assert_equal(I, 0.3 * self.OI)
+
+        # WMG=nx.MultiGraph(self.WG)
+        # WMG.add_edge(0,1,weight=0.5,other=0.3)
+        # np.testing.assert_equal(nx.incidence_matrix(WMG,weight='weight').todense(),
+        #              np.abs(0.5*self.MGOI))
+        # np.testing.assert_equal(nx.incidence_matrix(WMG,weight='weight',oriented=True).todense(),
+        #              0.5*self.MGOI)
+        # np.testing.assert_equal(nx.incidence_matrix(WMG,weight='other',oriented=True).todense(),
+        #              0.3*self.MGOI)
+
+        WMG = nx.MultiGraph(self.WG)
+        WMG.add_edge(0, 1, weight=0.5, other=0.3)
+
+        I = nx.incidence_matrix(
+            WMG,
+            nodelist=sorted(WMG),
+            edgelist=sorted(WMG.edges(keys=True)),
+            oriented=True,
+            weight="weight",
+        ).todense()
+        np.testing.assert_equal(I, 0.5 * self.MGOI)
+
+        I = nx.incidence_matrix(
+            WMG,
+            nodelist=sorted(WMG),
+            edgelist=sorted(WMG.edges(keys=True)),
+            oriented=False,
+            weight="weight",
+        ).todense()
+        np.testing.assert_equal(I, np.abs(0.5 * self.MGOI))
+
+        I = nx.incidence_matrix(
+            WMG,
+            nodelist=sorted(WMG),
+            edgelist=sorted(WMG.edges(keys=True)),
+            oriented=True,
+            weight="other",
+        ).todense()
+        np.testing.assert_equal(I, 0.3 * self.MGOI)
+
+    def test_adjacency_matrix(self):
+        "Conversion to adjacency matrix"
+        np.testing.assert_equal(nx.adjacency_matrix(self.G).todense(), self.A)
+        np.testing.assert_equal(nx.adjacency_matrix(self.MG).todense(), self.A)
+        np.testing.assert_equal(nx.adjacency_matrix(self.MG2).todense(), self.MG2A)
+        np.testing.assert_equal(
+            nx.adjacency_matrix(self.G, nodelist=[0, 1]).todense(), self.A[:2, :2]
+        )
+        np.testing.assert_equal(nx.adjacency_matrix(self.WG).todense(), self.WA)
+        np.testing.assert_equal(
+            nx.adjacency_matrix(self.WG, weight=None).todense(), self.A
+        )
+        np.testing.assert_equal(
+            nx.adjacency_matrix(self.MG2, weight=None).todense(), self.MG2A
+        )
+        np.testing.assert_equal(
+            nx.adjacency_matrix(self.WG, weight="other").todense(), 0.6 * self.WA
+        )
+        np.testing.assert_equal(
+            nx.adjacency_matrix(self.no_edges_G, nodelist=[1, 3]).todense(),
+            self.no_edges_A,
+        )

pythonProject/.venv/Lib/site-packages/networkx/linalg/tests/test_laplacian.py

@@ -0,0 +1,336 @@
+import pytest
+
+np = pytest.importorskip("numpy")
+pytest.importorskip("scipy")
+
+import networkx as nx
+from networkx.generators.degree_seq import havel_hakimi_graph
+from networkx.generators.expanders import margulis_gabber_galil_graph
+
+
+class TestLaplacian:
+    @classmethod
+    def setup_class(cls):
+        deg = [3, 2, 2, 1, 0]
+        cls.G = havel_hakimi_graph(deg)
+        cls.WG = nx.Graph(
+            (u, v, {"weight": 0.5, "other": 0.3}) for (u, v) in cls.G.edges()
+        )
+        cls.WG.add_node(4)
+        cls.MG = nx.MultiGraph(cls.G)
+
+        # Graph with clsloops
+        cls.Gsl = cls.G.copy()
+        for node in cls.Gsl.nodes():
+            cls.Gsl.add_edge(node, node)
+
+        # Graph used as an example in Sec. 4.1 of Langville and Meyer,
+        # "Google's PageRank and Beyond".
+        cls.DiG = nx.DiGraph()
+        cls.DiG.add_edges_from(
+            (
+                (1, 2),
+                (1, 3),
+                (3, 1),
+                (3, 2),
+                (3, 5),
+                (4, 5),
+                (4, 6),
+                (5, 4),
+                (5, 6),
+                (6, 4),
+            )
+        )
+        cls.DiMG = nx.MultiDiGraph(cls.DiG)
+        cls.DiWG = nx.DiGraph(
+            (u, v, {"weight": 0.5, "other": 0.3}) for (u, v) in cls.DiG.edges()
+        )
+        cls.DiGsl = cls.DiG.copy()
+        for node in cls.DiGsl.nodes():
+            cls.DiGsl.add_edge(node, node)
+
+    def test_laplacian(self):
+        "Graph Laplacian"
+        # fmt: off
+        NL = np.array([[ 3, -1, -1, -1,  0],
+                       [-1,  2, -1,  0,  0],
+                       [-1, -1,  2,  0,  0],
+                       [-1,  0,  0,  1,  0],
+                       [ 0,  0,  0,  0,  0]])
+        # fmt: on
+        WL = 0.5 * NL
+        OL = 0.3 * NL
+        # fmt: off
+        DiNL = np.array([[ 2, -1, -1,  0,  0,  0],
+                         [ 0,  0,  0,  0,  0,  0],
+                         [-1, -1,  3, -1,  0,  0],
+                         [ 0,  0,  0,  2, -1, -1],
+                         [ 0,  0,  0, -1,  2, -1],
+                         [ 0,  0,  0,  0, -1,  1]])
+        # fmt: on
+        DiWL = 0.5 * DiNL
+        DiOL = 0.3 * DiNL
+        np.testing.assert_equal(nx.laplacian_matrix(self.G).todense(), NL)
+        np.testing.assert_equal(nx.laplacian_matrix(self.MG).todense(), NL)
+        np.testing.assert_equal(
+            nx.laplacian_matrix(self.G, nodelist=[0, 1]).todense(),
+            np.array([[1, -1], [-1, 1]]),
+        )
+        np.testing.assert_equal(nx.laplacian_matrix(self.WG).todense(), WL)
+        np.testing.assert_equal(nx.laplacian_matrix(self.WG, weight=None).todense(), NL)
+        np.testing.assert_equal(
+            nx.laplacian_matrix(self.WG, weight="other").todense(), OL
+        )
+
+        np.testing.assert_equal(nx.laplacian_matrix(self.DiG).todense(), DiNL)
+        np.testing.assert_equal(nx.laplacian_matrix(self.DiMG).todense(), DiNL)
+        np.testing.assert_equal(
+            nx.laplacian_matrix(self.DiG, nodelist=[1, 2]).todense(),
+            np.array([[1, -1], [0, 0]]),
+        )
+        np.testing.assert_equal(nx.laplacian_matrix(self.DiWG).todense(), DiWL)
+        np.testing.assert_equal(
+            nx.laplacian_matrix(self.DiWG, weight=None).todense(), DiNL
+        )
+        np.testing.assert_equal(
+            nx.laplacian_matrix(self.DiWG, weight="other").todense(), DiOL
+        )
+
+    def test_normalized_laplacian(self):
+        "Generalized Graph Laplacian"
+        # fmt: off
+        G = np.array([[ 1.   , -0.408, -0.408, -0.577,  0.],
+                      [-0.408,  1.   , -0.5  ,  0.   ,  0.],
+                      [-0.408, -0.5  ,  1.   ,  0.   ,  0.],
+                      [-0.577,  0.   ,  0.   ,  1.   ,  0.],
+                      [ 0.   ,  0.   ,  0.   ,  0.   ,  0.]])
+        GL = np.array([[ 1.   , -0.408, -0.408, -0.577,  0.   ],
+                       [-0.408,  1.   , -0.5  ,  0.   ,  0.   ],
+                       [-0.408, -0.5  ,  1.   ,  0.   ,  0.   ],
+                       [-0.577,  0.   ,  0.   ,  1.   ,  0.   ],
+                       [ 0.   ,  0.   ,  0.   ,  0.   ,  0.   ]])
+        Lsl = np.array([[ 0.75  , -0.2887, -0.2887, -0.3536,  0.    ],
+                        [-0.2887,  0.6667, -0.3333,  0.    ,  0.    ],
+                        [-0.2887, -0.3333,  0.6667,  0.    ,  0.    ],
+                        [-0.3536,  0.    ,  0.    ,  0.5   ,  0.    ],
+                        [ 0.    ,  0.    ,  0.    ,  0.    ,  0.    ]])
+
+        DiG = np.array([[ 1.    ,  0.    , -0.4082,  0.    ,  0.    ,  0.    ],
+                        [ 0.    ,  0.    ,  0.    ,  0.    ,  0.    ,  0.    ],
+                        [-0.4082,  0.    ,  1.    ,  0.    , -0.4082,  0.    ],
+                        [ 0.    ,  0.    ,  0.    ,  1.    , -0.5   , -0.7071],
+                        [ 0.    ,  0.    ,  0.    , -0.5   ,  1.    , -0.7071],
+                        [ 0.    ,  0.    ,  0.    , -0.7071,  0.     , 1.    ]])
+        DiGL = np.array([[ 1.    ,  0.    , -0.4082,  0.    ,  0.    ,  0.    ],
+                         [ 0.    ,  0.    ,  0.    ,  0.    ,  0.    ,  0.    ],
+                         [-0.4082,  0.    ,  1.    , -0.4082,  0.    ,  0.    ],
+                         [ 0.    ,  0.    ,  0.    ,  1.    , -0.5   , -0.7071],
+                         [ 0.    ,  0.    ,  0.    , -0.5   ,  1.    , -0.7071],
+                         [ 0.    ,  0.    ,  0.    ,  0.    , -0.7071,  1.    ]])
+        DiLsl = np.array([[ 0.6667, -0.5774, -0.2887,  0.    ,  0.    ,  0.    ],
+                          [ 0.    ,  0.    ,  0.    ,  0.    ,  0.    ,  0.    ],
+                          [-0.2887, -0.5   ,  0.75  , -0.2887,  0.    ,  0.    ],
+                          [ 0.    ,  0.    ,  0.    ,  0.6667, -0.3333, -0.4082],
+                          [ 0.    ,  0.    ,  0.    , -0.3333,  0.6667, -0.4082],
+                          [ 0.    ,  0.    ,  0.    ,  0.    , -0.4082,  0.5   ]])
+        # fmt: on
+
+        np.testing.assert_almost_equal(
+            nx.normalized_laplacian_matrix(self.G, nodelist=range(5)).todense(),
+            G,
+            decimal=3,
+        )
+        np.testing.assert_almost_equal(
+            nx.normalized_laplacian_matrix(self.G).todense(), GL, decimal=3
+        )
+        np.testing.assert_almost_equal(
+            nx.normalized_laplacian_matrix(self.MG).todense(), GL, decimal=3
+        )
+        np.testing.assert_almost_equal(
+            nx.normalized_laplacian_matrix(self.WG).todense(), GL, decimal=3
+        )
+        np.testing.assert_almost_equal(
+            nx.normalized_laplacian_matrix(self.WG, weight="other").todense(),
+            GL,
+            decimal=3,
+        )
+        np.testing.assert_almost_equal(
+            nx.normalized_laplacian_matrix(self.Gsl).todense(), Lsl, decimal=3
+        )
+
+        np.testing.assert_almost_equal(
+            nx.normalized_laplacian_matrix(
+                self.DiG,
+                nodelist=range(1, 1 + 6),
+            ).todense(),
+            DiG,
+            decimal=3,
+        )
+        np.testing.assert_almost_equal(
+            nx.normalized_laplacian_matrix(self.DiG).todense(), DiGL, decimal=3
+        )
+        np.testing.assert_almost_equal(
+            nx.normalized_laplacian_matrix(self.DiMG).todense(), DiGL, decimal=3
+        )
+        np.testing.assert_almost_equal(
+            nx.normalized_laplacian_matrix(self.DiWG).todense(), DiGL, decimal=3
+        )
+        np.testing.assert_almost_equal(
+            nx.normalized_laplacian_matrix(self.DiWG, weight="other").todense(),
+            DiGL,
+            decimal=3,
+        )
+        np.testing.assert_almost_equal(
+            nx.normalized_laplacian_matrix(self.DiGsl).todense(), DiLsl, decimal=3
+        )
+
+
+def test_directed_laplacian():
+    "Directed Laplacian"
+    # Graph used as an example in Sec. 4.1 of Langville and Meyer,
+    # "Google's PageRank and Beyond". The graph contains dangling nodes, so
+    # the pagerank random walk is selected by directed_laplacian
+    G = nx.DiGraph()
+    G.add_edges_from(
+        (
+            (1, 2),
+            (1, 3),
+            (3, 1),
+            (3, 2),
+            (3, 5),
+            (4, 5),
+            (4, 6),
+            (5, 4),
+            (5, 6),
+            (6, 4),
+        )
+    )
+    # fmt: off
+    GL = np.array([[ 0.9833, -0.2941, -0.3882, -0.0291, -0.0231, -0.0261],
+                   [-0.2941,  0.8333, -0.2339, -0.0536, -0.0589, -0.0554],
+                   [-0.3882, -0.2339,  0.9833, -0.0278, -0.0896, -0.0251],
+                   [-0.0291, -0.0536, -0.0278,  0.9833, -0.4878, -0.6675],
+                   [-0.0231, -0.0589, -0.0896, -0.4878,  0.9833, -0.2078],
+                   [-0.0261, -0.0554, -0.0251, -0.6675, -0.2078,  0.9833]])
+    # fmt: on
+    L = nx.directed_laplacian_matrix(G, alpha=0.9, nodelist=sorted(G))
+    np.testing.assert_almost_equal(L, GL, decimal=3)
+
+    # Make the graph strongly connected, so we can use a random and lazy walk
+    G.add_edges_from(((2, 5), (6, 1)))
+    # fmt: off
+    GL = np.array([[ 1.    , -0.3062, -0.4714,  0.    ,  0.    , -0.3227],
+                   [-0.3062,  1.    , -0.1443,  0.    , -0.3162,  0.    ],
+                   [-0.4714, -0.1443,  1.    ,  0.    , -0.0913,  0.    ],
+                   [ 0.    ,  0.    ,  0.    ,  1.    , -0.5   , -0.5   ],
+                   [ 0.    , -0.3162, -0.0913, -0.5   ,  1.    , -0.25  ],
+                   [-0.3227,  0.    ,  0.    , -0.5   , -0.25  ,  1.    ]])
+    # fmt: on
+    L = nx.directed_laplacian_matrix(
+        G, alpha=0.9, nodelist=sorted(G), walk_type="random"
+    )
+    np.testing.assert_almost_equal(L, GL, decimal=3)
+
+    # fmt: off
+    GL = np.array([[ 0.5   , -0.1531, -0.2357,  0.    ,  0.    , -0.1614],
+                   [-0.1531,  0.5   , -0.0722,  0.    , -0.1581,  0.    ],
+                   [-0.2357, -0.0722,  0.5   ,  0.    , -0.0456,  0.    ],
+                   [ 0.    ,  0.    ,  0.    ,  0.5   , -0.25  , -0.25  ],
+                   [ 0.    , -0.1581, -0.0456, -0.25  ,  0.5   , -0.125 ],
+                   [-0.1614,  0.    ,  0.    , -0.25  , -0.125 ,  0.5   ]])  
+    # fmt: on
+    L = nx.directed_laplacian_matrix(G, alpha=0.9, nodelist=sorted(G), walk_type="lazy")
+    np.testing.assert_almost_equal(L, GL, decimal=3)
+
+    # Make a strongly connected periodic graph
+    G = nx.DiGraph()
+    G.add_edges_from(((1, 2), (2, 4), (4, 1), (1, 3), (3, 4)))
+    # fmt: off
+    GL = np.array([[ 0.5  , -0.176, -0.176, -0.25 ],
+                   [-0.176,  0.5  ,  0.   , -0.176],
+                   [-0.176,  0.   ,  0.5  , -0.176],
+                   [-0.25 , -0.176, -0.176,  0.5  ]])
+    # fmt: on
+    L = nx.directed_laplacian_matrix(G, alpha=0.9, nodelist=sorted(G))
+    np.testing.assert_almost_equal(L, GL, decimal=3)
+
+
+def test_directed_combinatorial_laplacian():
+    "Directed combinatorial Laplacian"
+    # Graph used as an example in Sec. 4.1 of Langville and Meyer,
+    # "Google's PageRank and Beyond". The graph contains dangling nodes, so
+    # the pagerank random walk is selected by directed_laplacian
+    G = nx.DiGraph()
+    G.add_edges_from(
+        (
+            (1, 2),
+            (1, 3),
+            (3, 1),
+            (3, 2),
+            (3, 5),
+            (4, 5),
+            (4, 6),
+            (5, 4),
+            (5, 6),
+            (6, 4),
+        )
+    )
+    # fmt: off
+    GL = np.array([[ 0.0366, -0.0132, -0.0153, -0.0034, -0.0020, -0.0027],
+                   [-0.0132,  0.0450, -0.0111, -0.0076, -0.0062, -0.0069],
+                   [-0.0153, -0.0111,  0.0408, -0.0035, -0.0083, -0.0027],
+                   [-0.0034, -0.0076, -0.0035,  0.3688, -0.1356, -0.2187],
+                   [-0.0020, -0.0062, -0.0083, -0.1356,  0.2026, -0.0505],
+                   [-0.0027, -0.0069, -0.0027, -0.2187, -0.0505,  0.2815]])
+    # fmt: on
+
+    L = nx.directed_combinatorial_laplacian_matrix(G, alpha=0.9, nodelist=sorted(G))
+    np.testing.assert_almost_equal(L, GL, decimal=3)
+
+    # Make the graph strongly connected, so we can use a random and lazy walk
+    G.add_edges_from(((2, 5), (6, 1)))
+
+    # fmt: off
+    GL = np.array([[ 0.1395, -0.0349, -0.0465,  0.    ,  0.    , -0.0581],
+                   [-0.0349,  0.093 , -0.0116,  0.    , -0.0465,  0.    ],
+                   [-0.0465, -0.0116,  0.0698,  0.    , -0.0116,  0.    ],
+                   [ 0.    ,  0.    ,  0.    ,  0.2326, -0.1163, -0.1163],
+                   [ 0.    , -0.0465, -0.0116, -0.1163,  0.2326, -0.0581],
+                   [-0.0581,  0.    ,  0.    , -0.1163, -0.0581,  0.2326]])
+    # fmt: on
+
+    L = nx.directed_combinatorial_laplacian_matrix(
+        G, alpha=0.9, nodelist=sorted(G), walk_type="random"
+    )
+    np.testing.assert_almost_equal(L, GL, decimal=3)
+
+    # fmt: off
+    GL = np.array([[ 0.0698, -0.0174, -0.0233,  0.    ,  0.    , -0.0291],
+                   [-0.0174,  0.0465, -0.0058,  0.    , -0.0233,  0.    ],
+                   [-0.0233, -0.0058,  0.0349,  0.    , -0.0058,  0.    ],
+                   [ 0.    ,  0.    ,  0.    ,  0.1163, -0.0581, -0.0581],
+                   [ 0.    , -0.0233, -0.0058, -0.0581,  0.1163, -0.0291],
+                   [-0.0291,  0.    ,  0.    , -0.0581, -0.0291,  0.1163]])
+    # fmt: on
+
+    L = nx.directed_combinatorial_laplacian_matrix(
+        G, alpha=0.9, nodelist=sorted(G), walk_type="lazy"
+    )
+    np.testing.assert_almost_equal(L, GL, decimal=3)
+
+    E = nx.DiGraph(margulis_gabber_galil_graph(2))
+    L = nx.directed_combinatorial_laplacian_matrix(E)
+    # fmt: off
+    expected = np.array(
+        [[ 0.16666667, -0.08333333, -0.08333333,  0.        ],
+         [-0.08333333,  0.16666667,  0.        , -0.08333333],
+         [-0.08333333,  0.        ,  0.16666667, -0.08333333],
+         [ 0.        , -0.08333333, -0.08333333,  0.16666667]]
+    )
+    # fmt: on
+    np.testing.assert_almost_equal(L, expected, decimal=6)
+
+    with pytest.raises(nx.NetworkXError):
+        nx.directed_combinatorial_laplacian_matrix(G, walk_type="pagerank", alpha=100)
+    with pytest.raises(nx.NetworkXError):
+        nx.directed_combinatorial_laplacian_matrix(G, walk_type="silly")

pythonProject/.venv/Lib/site-packages/networkx/linalg/tests/test_modularity.py

@@ -0,0 +1,87 @@
+import pytest
+
+np = pytest.importorskip("numpy")
+pytest.importorskip("scipy")
+
+import networkx as nx
+from networkx.generators.degree_seq import havel_hakimi_graph
+
+
+class TestModularity:
+    @classmethod
+    def setup_class(cls):
+        deg = [3, 2, 2, 1, 0]
+        cls.G = havel_hakimi_graph(deg)
+        # Graph used as an example in Sec. 4.1 of Langville and Meyer,
+        # "Google's PageRank and Beyond". (Used for test_directed_laplacian)
+        cls.DG = nx.DiGraph()
+        cls.DG.add_edges_from(
+            (
+                (1, 2),
+                (1, 3),
+                (3, 1),
+                (3, 2),
+                (3, 5),
+                (4, 5),
+                (4, 6),
+                (5, 4),
+                (5, 6),
+                (6, 4),
+            )
+        )
+
+    def test_modularity(self):
+        "Modularity matrix"
+        # fmt: off
+        B = np.array([[-1.125,  0.25,  0.25,  0.625,  0.],
+                      [0.25, -0.5,  0.5, -0.25,  0.],
+                      [0.25,  0.5, -0.5, -0.25,  0.],
+                      [0.625, -0.25, -0.25, -0.125,  0.],
+                      [0.,  0.,  0.,  0.,  0.]])
+        # fmt: on
+
+        permutation = [4, 0, 1, 2, 3]
+        np.testing.assert_equal(nx.modularity_matrix(self.G), B)
+        np.testing.assert_equal(
+            nx.modularity_matrix(self.G, nodelist=permutation),
+            B[np.ix_(permutation, permutation)],
+        )
+
+    def test_modularity_weight(self):
+        "Modularity matrix with weights"
+        # fmt: off
+        B = np.array([[-1.125,  0.25,  0.25,  0.625,  0.],
+                      [0.25, -0.5,  0.5, -0.25,  0.],
+                      [0.25,  0.5, -0.5, -0.25,  0.],
+                      [0.625, -0.25, -0.25, -0.125,  0.],
+                      [0.,  0.,  0.,  0.,  0.]])
+        # fmt: on
+
+        G_weighted = self.G.copy()
+        for n1, n2 in G_weighted.edges():
+            G_weighted.edges[n1, n2]["weight"] = 0.5
+        # The following test would fail in networkx 1.1
+        np.testing.assert_equal(nx.modularity_matrix(G_weighted), B)
+        # The following test that the modularity matrix get rescaled accordingly
+        np.testing.assert_equal(
+            nx.modularity_matrix(G_weighted, weight="weight"), 0.5 * B
+        )
+
+    def test_directed_modularity(self):
+        "Directed Modularity matrix"
+        # fmt: off
+        B = np.array([[-0.2,  0.6,  0.8, -0.4, -0.4, -0.4],
+                      [0.,  0.,  0.,  0.,  0.,  0.],
+                      [0.7,  0.4, -0.3, -0.6,  0.4, -0.6],
+                      [-0.2, -0.4, -0.2, -0.4,  0.6,  0.6],
+                      [-0.2, -0.4, -0.2,  0.6, -0.4,  0.6],
+                      [-0.1, -0.2, -0.1,  0.8, -0.2, -0.2]])
+        # fmt: on
+        node_permutation = [5, 1, 2, 3, 4, 6]
+        idx_permutation = [4, 0, 1, 2, 3, 5]
+        mm = nx.directed_modularity_matrix(self.DG, nodelist=sorted(self.DG))
+        np.testing.assert_equal(mm, B)
+        np.testing.assert_equal(
+            nx.directed_modularity_matrix(self.DG, nodelist=node_permutation),
+            B[np.ix_(idx_permutation, idx_permutation)],
+        )

pythonProject/.venv/Lib/site-packages/networkx/linalg/tests/test_spectrum.py

@@ -0,0 +1,71 @@
+import pytest
+
+np = pytest.importorskip("numpy")
+pytest.importorskip("scipy")
+
+import networkx as nx
+from networkx.generators.degree_seq import havel_hakimi_graph
+
+
+class TestSpectrum:
+    @classmethod
+    def setup_class(cls):
+        deg = [3, 2, 2, 1, 0]
+        cls.G = havel_hakimi_graph(deg)
+        cls.P = nx.path_graph(3)
+        cls.WG = nx.Graph(
+            (u, v, {"weight": 0.5, "other": 0.3}) for (u, v) in cls.G.edges()
+        )
+        cls.WG.add_node(4)
+        cls.DG = nx.DiGraph()
+        nx.add_path(cls.DG, [0, 1, 2])
+
+    def test_laplacian_spectrum(self):
+        "Laplacian eigenvalues"
+        evals = np.array([0, 0, 1, 3, 4])
+        e = sorted(nx.laplacian_spectrum(self.G))
+        np.testing.assert_almost_equal(e, evals)
+        e = sorted(nx.laplacian_spectrum(self.WG, weight=None))
+        np.testing.assert_almost_equal(e, evals)
+        e = sorted(nx.laplacian_spectrum(self.WG))
+        np.testing.assert_almost_equal(e, 0.5 * evals)
+        e = sorted(nx.laplacian_spectrum(self.WG, weight="other"))
+        np.testing.assert_almost_equal(e, 0.3 * evals)
+
+    def test_normalized_laplacian_spectrum(self):
+        "Normalized Laplacian eigenvalues"
+        evals = np.array([0, 0, 0.7712864461218, 1.5, 1.7287135538781])
+        e = sorted(nx.normalized_laplacian_spectrum(self.G))
+        np.testing.assert_almost_equal(e, evals)
+        e = sorted(nx.normalized_laplacian_spectrum(self.WG, weight=None))
+        np.testing.assert_almost_equal(e, evals)
+        e = sorted(nx.normalized_laplacian_spectrum(self.WG))
+        np.testing.assert_almost_equal(e, evals)
+        e = sorted(nx.normalized_laplacian_spectrum(self.WG, weight="other"))
+        np.testing.assert_almost_equal(e, evals)
+
+    def test_adjacency_spectrum(self):
+        "Adjacency eigenvalues"
+        evals = np.array([-np.sqrt(2), 0, np.sqrt(2)])
+        e = sorted(nx.adjacency_spectrum(self.P))
+        np.testing.assert_almost_equal(e, evals)
+
+    def test_modularity_spectrum(self):
+        "Modularity eigenvalues"
+        evals = np.array([-1.5, 0.0, 0.0])
+        e = sorted(nx.modularity_spectrum(self.P))
+        np.testing.assert_almost_equal(e, evals)
+        # Directed modularity eigenvalues
+        evals = np.array([-0.5, 0.0, 0.0])
+        e = sorted(nx.modularity_spectrum(self.DG))
+        np.testing.assert_almost_equal(e, evals)
+
+    def test_bethe_hessian_spectrum(self):
+        "Bethe Hessian eigenvalues"
+        evals = np.array([0.5 * (9 - np.sqrt(33)), 4, 0.5 * (9 + np.sqrt(33))])
+        e = sorted(nx.bethe_hessian_spectrum(self.P, r=2))
+        np.testing.assert_almost_equal(e, evals)
+        # Collapses back to Laplacian:
+        e1 = sorted(nx.bethe_hessian_spectrum(self.P, r=1))
+        e2 = sorted(nx.laplacian_spectrum(self.P))
+        np.testing.assert_almost_equal(e1, e2)

pythonProject/.venv/Lib/site-packages/numpy/ctypeslib.py

@@ -0,0 +1,602 @@
+"""
+============================
+``ctypes`` Utility Functions
+============================
+
+See Also
+--------
+load_library : Load a C library.
+ndpointer : Array restype/argtype with verification.
+as_ctypes : Create a ctypes array from an ndarray.
+as_array : Create an ndarray from a ctypes array.
+
+References
+----------
+.. [1] "SciPy Cookbook: ctypes", https://scipy-cookbook.readthedocs.io/items/Ctypes.html
+
+Examples
+--------
+Load the C library:
+
+>>> _lib = np.ctypeslib.load_library('libmystuff', '.')     #doctest: +SKIP
+
+Our result type, an ndarray that must be of type double, be 1-dimensional
+and is C-contiguous in memory:
+
+>>> array_1d_double = np.ctypeslib.ndpointer(
+...                          dtype=np.double,
+...                          ndim=1, flags='CONTIGUOUS')    #doctest: +SKIP
+
+Our C-function typically takes an array and updates its values
+in-place.  For example::
+
+    void foo_func(double* x, int length)
+    {
+        int i;
+        for (i = 0; i < length; i++) {
+            x[i] = i*i;
+        }
+    }
+
+We wrap it using:
+
+>>> _lib.foo_func.restype = None                      #doctest: +SKIP
+>>> _lib.foo_func.argtypes = [array_1d_double, c_int] #doctest: +SKIP
+
+Then, we're ready to call ``foo_func``:
+
+>>> out = np.empty(15, dtype=np.double)
+>>> _lib.foo_func(out, len(out))                #doctest: +SKIP
+
+"""
+__all__ = ['load_library', 'ndpointer', 'c_intp', 'as_ctypes', 'as_array',
+           'as_ctypes_type']
+
+import os
+import numpy as np
+from numpy._core.multiarray import _flagdict, flagsobj
+
+try:
+    import ctypes
+except ImportError:
+    ctypes = None
+
+if ctypes is None:
+    def _dummy(*args, **kwds):
+        """
+        Dummy object that raises an ImportError if ctypes is not available.
+
+        Raises
+        ------
+        ImportError
+            If ctypes is not available.
+
+        """
+        raise ImportError("ctypes is not available.")
+    load_library = _dummy
+    as_ctypes = _dummy
+    as_array = _dummy
+    from numpy import intp as c_intp
+    _ndptr_base = object
+else:
+    import numpy._core._internal as nic
+    c_intp = nic._getintp_ctype()
+    del nic
+    _ndptr_base = ctypes.c_void_p
+
+    # Adapted from Albert Strasheim
+    def load_library(libname, loader_path):
+        """
+        It is possible to load a library using
+
+        >>> lib = ctypes.cdll[<full_path_name>] # doctest: +SKIP
+
+        But there are cross-platform considerations, such as library file extensions,
+        plus the fact Windows will just load the first library it finds with that name.
+        NumPy supplies the load_library function as a convenience.
+
+        .. versionchanged:: 1.20.0
+            Allow libname and loader_path to take any
+            :term:`python:path-like object`.
+
+        Parameters
+        ----------
+        libname : path-like
+            Name of the library, which can have 'lib' as a prefix,
+            but without an extension.
+        loader_path : path-like
+            Where the library can be found.
+
+        Returns
+        -------
+        ctypes.cdll[libpath] : library object
+           A ctypes library object
+
+        Raises
+        ------
+        OSError
+            If there is no library with the expected extension, or the
+            library is defective and cannot be loaded.
+        """
+        # Convert path-like objects into strings
+        libname = os.fsdecode(libname)
+        loader_path = os.fsdecode(loader_path)
+
+        ext = os.path.splitext(libname)[1]
+        if not ext:
+            import sys
+            import sysconfig
+            # Try to load library with platform-specific name, otherwise
+            # default to libname.[so|dll|dylib].  Sometimes, these files are
+            # built erroneously on non-linux platforms.
+            base_ext = ".so"
+            if sys.platform.startswith("darwin"):
+                base_ext = ".dylib"
+            elif sys.platform.startswith("win"):
+                base_ext = ".dll"
+            libname_ext = [libname + base_ext]
+            so_ext = sysconfig.get_config_var("EXT_SUFFIX")
+            if not so_ext == base_ext:
+                libname_ext.insert(0, libname + so_ext)
+        else:
+            libname_ext = [libname]
+
+        loader_path = os.path.abspath(loader_path)
+        if not os.path.isdir(loader_path):
+            libdir = os.path.dirname(loader_path)
+        else:
+            libdir = loader_path
+
+        for ln in libname_ext:
+            libpath = os.path.join(libdir, ln)
+            if os.path.exists(libpath):
+                try:
+                    return ctypes.cdll[libpath]
+                except OSError:
+                    ## defective lib file
+                    raise
+        ## if no successful return in the libname_ext loop:
+        raise OSError("no file with expected extension")
+
+
+def _num_fromflags(flaglist):
+    num = 0
+    for val in flaglist:
+        num += _flagdict[val]
+    return num
+
+_flagnames = ['C_CONTIGUOUS', 'F_CONTIGUOUS', 'ALIGNED', 'WRITEABLE',
+              'OWNDATA', 'WRITEBACKIFCOPY']
+def _flags_fromnum(num):
+    res = []
+    for key in _flagnames:
+        value = _flagdict[key]
+        if (num & value):
+            res.append(key)
+    return res
+
+
+class _ndptr(_ndptr_base):
+    @classmethod
+    def from_param(cls, obj):
+        if not isinstance(obj, np.ndarray):
+            raise TypeError("argument must be an ndarray")
+        if cls._dtype_ is not None \
+               and obj.dtype != cls._dtype_:
+            raise TypeError("array must have data type %s" % cls._dtype_)
+        if cls._ndim_ is not None \
+               and obj.ndim != cls._ndim_:
+            raise TypeError("array must have %d dimension(s)" % cls._ndim_)
+        if cls._shape_ is not None \
+               and obj.shape != cls._shape_:
+            raise TypeError("array must have shape %s" % str(cls._shape_))
+        if cls._flags_ is not None \
+               and ((obj.flags.num & cls._flags_) != cls._flags_):
+            raise TypeError("array must have flags %s" %
+                    _flags_fromnum(cls._flags_))
+        return obj.ctypes
+
+
+class _concrete_ndptr(_ndptr):
+    """
+    Like _ndptr, but with `_shape_` and `_dtype_` specified.
+
+    Notably, this means the pointer has enough information to reconstruct
+    the array, which is not generally true.
+    """
+    def _check_retval_(self):
+        """
+        This method is called when this class is used as the .restype
+        attribute for a shared-library function, to automatically wrap the
+        pointer into an array.
+        """
+        return self.contents
+
+    @property
+    def contents(self):
+        """
+        Get an ndarray viewing the data pointed to by this pointer.
+
+        This mirrors the `contents` attribute of a normal ctypes pointer
+        """
+        full_dtype = np.dtype((self._dtype_, self._shape_))
+        full_ctype = ctypes.c_char * full_dtype.itemsize
+        buffer = ctypes.cast(self, ctypes.POINTER(full_ctype)).contents
+        return np.frombuffer(buffer, dtype=full_dtype).squeeze(axis=0)
+
+
+# Factory for an array-checking class with from_param defined for
+#  use with ctypes argtypes mechanism
+_pointer_type_cache = {}
+def ndpointer(dtype=None, ndim=None, shape=None, flags=None):
+    """
+    Array-checking restype/argtypes.
+
+    An ndpointer instance is used to describe an ndarray in restypes
+    and argtypes specifications.  This approach is more flexible than
+    using, for example, ``POINTER(c_double)``, since several restrictions
+    can be specified, which are verified upon calling the ctypes function.
+    These include data type, number of dimensions, shape and flags.  If a
+    given array does not satisfy the specified restrictions,
+    a ``TypeError`` is raised.
+
+    Parameters
+    ----------
+    dtype : data-type, optional
+        Array data-type.
+    ndim : int, optional
+        Number of array dimensions.
+    shape : tuple of ints, optional
+        Array shape.
+    flags : str or tuple of str
+        Array flags; may be one or more of:
+
+        - C_CONTIGUOUS / C / CONTIGUOUS
+        - F_CONTIGUOUS / F / FORTRAN
+        - OWNDATA / O
+        - WRITEABLE / W
+        - ALIGNED / A
+        - WRITEBACKIFCOPY / X
+
+    Returns
+    -------
+    klass : ndpointer type object
+        A type object, which is an ``_ndtpr`` instance containing
+        dtype, ndim, shape and flags information.
+
+    Raises
+    ------
+    TypeError
+        If a given array does not satisfy the specified restrictions.
+
+    Examples
+    --------
+    >>> clib.somefunc.argtypes = [np.ctypeslib.ndpointer(dtype=np.float64,
+    ...                                                  ndim=1,
+    ...                                                  flags='C_CONTIGUOUS')]
+    ... #doctest: +SKIP
+    >>> clib.somefunc(np.array([1, 2, 3], dtype=np.float64))
+    ... #doctest: +SKIP
+
+    """
+
+    # normalize dtype to dtype | None
+    if dtype is not None:
+        dtype = np.dtype(dtype)
+
+    # normalize flags to int | None
+    num = None
+    if flags is not None:
+        if isinstance(flags, str):
+            flags = flags.split(',')
+        elif isinstance(flags, (int, np.integer)):
+            num = flags
+            flags = _flags_fromnum(num)
+        elif isinstance(flags, flagsobj):
+            num = flags.num
+            flags = _flags_fromnum(num)
+        if num is None:
+            try:
+                flags = [x.strip().upper() for x in flags]
+            except Exception as e:
+                raise TypeError("invalid flags specification") from e
+            num = _num_fromflags(flags)
+
+    # normalize shape to tuple | None
+    if shape is not None:
+        try:
+            shape = tuple(shape)
+        except TypeError:
+            # single integer -> 1-tuple
+            shape = (shape,)
+
+    cache_key = (dtype, ndim, shape, num)
+
+    try:
+        return _pointer_type_cache[cache_key]
+    except KeyError:
+        pass
+
+    # produce a name for the new type
+    if dtype is None:
+        name = 'any'
+    elif dtype.names is not None:
+        name = str(id(dtype))
+    else:
+        name = dtype.str
+    if ndim is not None:
+        name += "_%dd" % ndim
+    if shape is not None:
+        name += "_"+"x".join(str(x) for x in shape)
+    if flags is not None:
+        name += "_"+"_".join(flags)
+
+    if dtype is not None and shape is not None:
+        base = _concrete_ndptr
+    else:
+        base = _ndptr
+
+    klass = type("ndpointer_%s"%name, (base,),
+                 {"_dtype_": dtype,
+                  "_shape_" : shape,
+                  "_ndim_" : ndim,
+                  "_flags_" : num})
+    _pointer_type_cache[cache_key] = klass
+    return klass
+
+
+if ctypes is not None:
+    def _ctype_ndarray(element_type, shape):
+        """ Create an ndarray of the given element type and shape """
+        for dim in shape[::-1]:
+            element_type = dim * element_type
+            # prevent the type name include np.ctypeslib
+            element_type.__module__ = None
+        return element_type
+
+
+    def _get_scalar_type_map():
+        """
+        Return a dictionary mapping native endian scalar dtype to ctypes types
+        """
+        ct = ctypes
+        simple_types = [
+            ct.c_byte, ct.c_short, ct.c_int, ct.c_long, ct.c_longlong,
+            ct.c_ubyte, ct.c_ushort, ct.c_uint, ct.c_ulong, ct.c_ulonglong,
+            ct.c_float, ct.c_double,
+            ct.c_bool,
+        ]
+        return {np.dtype(ctype): ctype for ctype in simple_types}
+
+
+    _scalar_type_map = _get_scalar_type_map()
+
+
+    def _ctype_from_dtype_scalar(dtype):
+        # swapping twice ensure that `=` is promoted to <, >, or |
+        dtype_with_endian = dtype.newbyteorder('S').newbyteorder('S')
+        dtype_native = dtype.newbyteorder('=')
+        try:
+            ctype = _scalar_type_map[dtype_native]
+        except KeyError as e:
+            raise NotImplementedError(
+                "Converting {!r} to a ctypes type".format(dtype)
+            ) from None
+
+        if dtype_with_endian.byteorder == '>':
+            ctype = ctype.__ctype_be__
+        elif dtype_with_endian.byteorder == '<':
+            ctype = ctype.__ctype_le__
+
+        return ctype
+
+
+    def _ctype_from_dtype_subarray(dtype):
+        element_dtype, shape = dtype.subdtype
+        ctype = _ctype_from_dtype(element_dtype)
+        return _ctype_ndarray(ctype, shape)
+
+
+    def _ctype_from_dtype_structured(dtype):
+        # extract offsets of each field
+        field_data = []
+        for name in dtype.names:
+            field_dtype, offset = dtype.fields[name][:2]
+            field_data.append((offset, name, _ctype_from_dtype(field_dtype)))
+
+        # ctypes doesn't care about field order
+        field_data = sorted(field_data, key=lambda f: f[0])
+
+        if len(field_data) > 1 and all(offset == 0 for offset, name, ctype in field_data):
+            # union, if multiple fields all at address 0
+            size = 0
+            _fields_ = []
+            for offset, name, ctype in field_data:
+                _fields_.append((name, ctype))
+                size = max(size, ctypes.sizeof(ctype))
+
+            # pad to the right size
+            if dtype.itemsize != size:
+                _fields_.append(('', ctypes.c_char * dtype.itemsize))
+
+            # we inserted manual padding, so always `_pack_`
+            return type('union', (ctypes.Union,), dict(
+                _fields_=_fields_,
+                _pack_=1,
+                __module__=None,
+            ))
+        else:
+            last_offset = 0
+            _fields_ = []
+            for offset, name, ctype in field_data:
+                padding = offset - last_offset
+                if padding < 0:
+                    raise NotImplementedError("Overlapping fields")
+                if padding > 0:
+                    _fields_.append(('', ctypes.c_char * padding))
+
+                _fields_.append((name, ctype))
+                last_offset = offset + ctypes.sizeof(ctype)
+
+
+            padding = dtype.itemsize - last_offset
+            if padding > 0:
+                _fields_.append(('', ctypes.c_char * padding))
+
+            # we inserted manual padding, so always `_pack_`
+            return type('struct', (ctypes.Structure,), dict(
+                _fields_=_fields_,
+                _pack_=1,
+                __module__=None,
+            ))
+
+
+    def _ctype_from_dtype(dtype):
+        if dtype.fields is not None:
+            return _ctype_from_dtype_structured(dtype)
+        elif dtype.subdtype is not None:
+            return _ctype_from_dtype_subarray(dtype)
+        else:
+            return _ctype_from_dtype_scalar(dtype)
+
+
+    def as_ctypes_type(dtype):
+        r"""
+        Convert a dtype into a ctypes type.
+
+        Parameters
+        ----------
+        dtype : dtype
+            The dtype to convert
+
+        Returns
+        -------
+        ctype
+            A ctype scalar, union, array, or struct
+
+        Raises
+        ------
+        NotImplementedError
+            If the conversion is not possible
+
+        Notes
+        -----
+        This function does not losslessly round-trip in either direction.
+
+        ``np.dtype(as_ctypes_type(dt))`` will:
+
+        - insert padding fields
+        - reorder fields to be sorted by offset
+        - discard field titles
+
+        ``as_ctypes_type(np.dtype(ctype))`` will:
+
+        - discard the class names of `ctypes.Structure`\ s and
+          `ctypes.Union`\ s
+        - convert single-element `ctypes.Union`\ s into single-element
+          `ctypes.Structure`\ s
+        - insert padding fields
+
+        Examples
+        --------
+        Converting a simple dtype:
+
+        >>> dt = np.dtype('int8')
+        >>> ctype = np.ctypeslib.as_ctypes_type(dt)
+        >>> ctype
+        <class 'ctypes.c_byte'>
+
+        Converting a structured dtype:
+
+        >>> dt = np.dtype([('x', 'i4'), ('y', 'f4')])
+        >>> ctype = np.ctypeslib.as_ctypes_type(dt)
+        >>> ctype
+        <class 'struct'>
+
+        """
+        return _ctype_from_dtype(np.dtype(dtype))
+
+
+    def as_array(obj, shape=None):
+        """
+        Create a numpy array from a ctypes array or POINTER.
+
+        The numpy array shares the memory with the ctypes object.
+
+        The shape parameter must be given if converting from a ctypes POINTER.
+        The shape parameter is ignored if converting from a ctypes array
+
+        Examples
+        --------
+        Converting a ctypes integer array:
+
+        >>> import ctypes
+        >>> ctypes_array = (ctypes.c_int * 5)(0, 1, 2, 3, 4)
+        >>> np_array = np.ctypeslib.as_array(ctypes_array)
+        >>> np_array
+        array([0, 1, 2, 3, 4], dtype=int32)
+
+        Converting a ctypes POINTER:
+
+        >>> import ctypes
+        >>> buffer = (ctypes.c_int * 5)(0, 1, 2, 3, 4)
+        >>> pointer = ctypes.cast(buffer, ctypes.POINTER(ctypes.c_int))
+        >>> np_array = np.ctypeslib.as_array(pointer, (5,))
+        >>> np_array
+        array([0, 1, 2, 3, 4], dtype=int32)
+
+        """
+        if isinstance(obj, ctypes._Pointer):
+            # convert pointers to an array of the desired shape
+            if shape is None:
+                raise TypeError(
+                    'as_array() requires a shape argument when called on a '
+                    'pointer')
+            p_arr_type = ctypes.POINTER(_ctype_ndarray(obj._type_, shape))
+            obj = ctypes.cast(obj, p_arr_type).contents
+
+        return np.asarray(obj)
+
+
+    def as_ctypes(obj):
+        """
+        Create and return a ctypes object from a numpy array.  Actually
+        anything that exposes the __array_interface__ is accepted.
+
+        Examples
+        --------
+        Create ctypes object from inferred int ``np.array``:
+
+        >>> inferred_int_array = np.array([1, 2, 3])
+        >>> c_int_array = np.ctypeslib.as_ctypes(inferred_int_array)
+        >>> type(c_int_array)
+        <class 'c_long_Array_3'>
+        >>> c_int_array[:]
+        [1, 2, 3]
+
+        Create ctypes object from explicit 8 bit unsigned int ``np.array`` :
+
+        >>> exp_int_array = np.array([1, 2, 3], dtype=np.uint8)
+        >>> c_int_array = np.ctypeslib.as_ctypes(exp_int_array)
+        >>> type(c_int_array)
+        <class 'c_ubyte_Array_3'>
+        >>> c_int_array[:]
+        [1, 2, 3]
+
+        """
+        ai = obj.__array_interface__
+        if ai["strides"]:
+            raise TypeError("strided arrays not supported")
+        if ai["version"] != 3:
+            raise TypeError("only __array_interface__ version 3 supported")
+        addr, readonly = ai["data"]
+        if readonly:
+            raise TypeError("readonly arrays unsupported")
+
+        # can't use `_dtype((ai["typestr"], ai["shape"]))` here, as it overflows
+        # dtype.itemsize (gh-14214)
+        ctype_scalar = as_ctypes_type(ai["typestr"])
+        result_type = _ctype_ndarray(ctype_scalar, ai["shape"])
+        result = result_type.from_address(addr)
+        result.__keep = obj
+        return result

pythonProject/.venv/Lib/site-packages/torch/include/ATen/detail/MTIAHooksInterface.h

@@ -0,0 +1,103 @@
+#pragma once
+
+#include <c10/core/Device.h>
+#include <c10/util/Exception.h>
+
+#include <c10/core/Stream.h>
+#include <c10/util/Registry.h>
+
+#include <ATen/detail/AcceleratorHooksInterface.h>
+
+#include <string>
+
+namespace at {
+class Context;
+}
+
+namespace at {
+
+constexpr const char* MTIA_HELP =
+    "The MTIA backend requires MTIA extension for PyTorch;"
+    "this error has occurred because you are trying "
+    "to use some MTIA's functionality without MTIA extension included.";
+
+struct TORCH_API MTIAHooksInterface : AcceleratorHooksInterface {
+// this fails the implementation if MTIAHooks functions are called, but
+// MTIA backend is not present.
+#define FAIL_MTIAHOOKS_FUNC(func) \
+  TORCH_CHECK(false, "Cannot execute ", func, "() without MTIA backend.");
+
+  ~MTIAHooksInterface() override = default;
+
+  virtual void initMTIA() const {
+    // Avoid logging here, since MTIA needs init devices first then it will know
+    // how many devices are available. Make it as no-op if mtia extension is not
+    // dynamically loaded.
+    return;
+  }
+
+  virtual bool hasMTIA() const {
+    return false;
+  }
+
+  DeviceIndex deviceCount() const override {
+    return 0;
+  }
+
+  virtual void deviceSynchronize(c10::DeviceIndex device_index) const {
+    FAIL_MTIAHOOKS_FUNC(__func__);
+  }
+
+  virtual std::string showConfig() const {
+    FAIL_MTIAHOOKS_FUNC(__func__);
+  }
+
+  bool hasPrimaryContext(DeviceIndex device_index) const override {
+    return false;
+  }
+
+  void setCurrentDevice(DeviceIndex device) const override {
+    FAIL_MTIAHOOKS_FUNC(__func__);
+  }
+
+  DeviceIndex getCurrentDevice() const override {
+    FAIL_MTIAHOOKS_FUNC(__func__);
+    return -1;
+  }
+
+  DeviceIndex exchangeDevice(DeviceIndex device) const override {
+    FAIL_MTIAHOOKS_FUNC(__func__);
+    return -1;
+  }
+
+  DeviceIndex maybeExchangeDevice(DeviceIndex device) const override {
+    FAIL_MTIAHOOKS_FUNC(__func__);
+    return -1;
+  }
+
+  virtual c10::Stream getCurrentStream(DeviceIndex device) const {
+    FAIL_MTIAHOOKS_FUNC(__func__);
+    return c10::Stream::unpack3(-1, 0, c10::DeviceType::MTIA);
+  }
+
+  virtual c10::Stream getDefaultStream(DeviceIndex device) const {
+    FAIL_MTIAHOOKS_FUNC(__func__);
+    return c10::Stream::unpack3(-1, 0, c10::DeviceType::MTIA);
+  }
+
+  virtual void setCurrentStream(const c10::Stream& stream) const {
+    FAIL_MTIAHOOKS_FUNC(__func__);
+  }
+};
+
+struct TORCH_API MTIAHooksArgs {};
+
+C10_DECLARE_REGISTRY(MTIAHooksRegistry, MTIAHooksInterface, MTIAHooksArgs);
+#define REGISTER_MTIA_HOOKS(clsname) \
+  C10_REGISTER_CLASS(MTIAHooksRegistry, clsname, clsname)
+
+namespace detail {
+TORCH_API const MTIAHooksInterface& getMTIAHooks();
+TORCH_API bool isMTIAHooksBuilt();
+} // namespace detail
+} // namespace at

pythonProject/.venv/Lib/site-packages/torch/include/ATen/detail/PrivateUse1HooksInterface.h

@@ -0,0 +1,61 @@
+#pragma once
+
+#include <ATen/core/Generator.h>
+#include <ATen/detail/AcceleratorHooksInterface.h>
+#include <c10/core/Allocator.h>
+#include <c10/core/Device.h>
+#include <c10/core/Storage.h>
+#include <c10/util/Exception.h>
+namespace at {
+
+struct TORCH_API PrivateUse1HooksInterface : AcceleratorHooksInterface {
+  ~PrivateUse1HooksInterface() override = default;
+  virtual const at::Generator& getDefaultGenerator(
+      c10::DeviceIndex device_index) {
+    TORCH_CHECK_NOT_IMPLEMENTED(
+        false,
+        "You should register `PrivateUse1HooksInterface` for PrivateUse1 before call `getDefaultGenerator`.");
+  }
+
+  virtual at::Device getDeviceFromPtr(void* data) const {
+    TORCH_CHECK_NOT_IMPLEMENTED(
+        false,
+        "You should register `PrivateUse1HooksInterface` for PrivateUse1 before call `getDeviceFromPtr`.");
+  }
+
+  virtual Allocator* getPinnedMemoryAllocator() const {
+    TORCH_CHECK(
+        false,
+        "You should register `PrivateUse1HooksInterface` for PrivateUse1 before call `getPinnedMemoryAllocator`.");
+  }
+
+  bool hasPrimaryContext(DeviceIndex device_index) const override {
+    TORCH_CHECK_NOT_IMPLEMENTED(
+        false,
+        "You should register `PrivateUse1HooksInterface` for PrivateUse1 before call `hasPrimaryContext`.");
+  }
+
+  virtual void initPrivateUse1() const {}
+  virtual void resizePrivateUse1Bytes(const c10::Storage &storage, size_t newsize) const {
+    TORCH_CHECK_NOT_IMPLEMENTED(
+        false,
+        "You should register `PrivateUse1HooksInterface` for PrivateUse1 before call `resizePrivateUse1Bytes`.");
+  }
+};
+
+struct TORCH_API PrivateUse1HooksArgs {};
+
+TORCH_API void RegisterPrivateUse1HooksInterface(
+    at::PrivateUse1HooksInterface* hook_);
+
+TORCH_API at::PrivateUse1HooksInterface* GetPrivateUse1HooksInterface();
+
+TORCH_API bool isPrivateUse1HooksRegistered();
+
+namespace detail {
+
+TORCH_API const at::PrivateUse1HooksInterface& getPrivateUse1Hooks();
+
+} // namespace detail
+
+} // namespace at

pythonProject/.venv/Lib/site-packages/torch/include/ATen/detail/XPUHooksInterface.h

@@ -0,0 +1,84 @@
+#pragma once
+
+#include <c10/core/Device.h>
+#include <c10/util/Exception.h>
+#include <ATen/core/Generator.h>
+#include <c10/util/Registry.h>
+
+namespace at {
+
+constexpr const char* XPU_HELP =
+    "The XPU backend requires Intel Extension for Pytorch;"
+    "this error has occurred because you are trying "
+    "to use some XPU's functionality, but the Intel Extension for Pytorch has not been "
+    "loaded for some reason. The Intel Extension for Pytorch MUST "
+    "be loaded, EVEN IF you don't directly use any symbols from that!";
+
+struct TORCH_API XPUHooksInterface {
+  virtual ~XPUHooksInterface() = default;
+
+  virtual void initXPU() const {
+    TORCH_CHECK(
+        false,
+        "Cannot initialize XPU without Intel Extension for Pytorch.",
+        XPU_HELP);
+  }
+
+  virtual bool hasXPU() const {
+    return false;
+  }
+
+  virtual std::string showConfig() const {
+    TORCH_CHECK(
+        false,
+        "Cannot query detailed XPU version without Intel Extension for Pytorch. ",
+        XPU_HELP);
+  }
+
+  virtual int32_t getGlobalIdxFromDevice(const Device& device) const {
+    TORCH_CHECK(false, "Cannot get XPU global device index without ATen_xpu library.");
+  }
+
+  virtual Generator getXPUGenerator(C10_UNUSED DeviceIndex device_index = -1) const {
+    TORCH_CHECK(false, "Cannot get XPU generator without Intel Extension for Pytorch. ", XPU_HELP);
+  }
+
+  virtual const Generator& getDefaultXPUGenerator(C10_UNUSED DeviceIndex device_index = -1) const {
+    TORCH_CHECK(false, "Cannot get default XPU generator without Intel Extension for Pytorch. ", XPU_HELP);
+  }
+
+  virtual DeviceIndex getNumGPUs() const {
+    return 0;
+  }
+
+  virtual DeviceIndex current_device() const {
+    TORCH_CHECK(false, "Cannot get current device on XPU without ATen_xpu library.");
+  }
+
+  virtual Device getDeviceFromPtr(void* /*data*/) const {
+    TORCH_CHECK(false, "Cannot get device of pointer on XPU without ATen_xpu library.");
+  }
+
+  virtual void deviceSynchronize(DeviceIndex /*device_index*/) const {
+    TORCH_CHECK(false, "Cannot synchronize XPU device without ATen_xpu library.");
+  }
+
+  virtual Allocator* getPinnedMemoryAllocator() const  {
+    TORCH_CHECK(false, "Cannot get XPU pinned memory allocator without ATen_xpu library.");
+  }
+
+  virtual bool isPinnedPtr(const void* /*data*/) const {
+    return false;
+  }
+};
+
+struct TORCH_API XPUHooksArgs {};
+
+C10_DECLARE_REGISTRY(XPUHooksRegistry, XPUHooksInterface, XPUHooksArgs);
+#define REGISTER_XPU_HOOKS(clsname) \
+  C10_REGISTER_CLASS(XPUHooksRegistry, clsname, clsname)
+
+namespace detail {
+TORCH_API const XPUHooksInterface& getXPUHooks();
+} // namespace detail
+} // namespace at

pythonProject/.venv/Lib/site-packages/torch/include/ATen/functorch/ADInterpreters.h

@@ -0,0 +1,38 @@
+#pragma once
+#include <ATen/functorch/Interpreter.h>
+
+namespace at::functorch {
+
+// These are the interpreters for our AD transforms
+// (grad, vjp and jvp).
+// See NOTE: [functorch interpreter stack] for more details.
+
+struct TORCH_API GradInterpreterPtr {
+  explicit GradInterpreterPtr(const Interpreter* base): base_(base) { TORCH_INTERNAL_ASSERT(base->key() == TransformType::Grad); }
+  TransformType key() const { return base_->key(); }
+  int64_t level() const { return base_->level(); }
+  void processImpl(const c10::OperatorHandle& op, torch::jit::Stack* stack);
+  void sendToNextInterpreterImpl(const c10::OperatorHandle& op, torch::jit::Stack* stack, bool grad_special_case);
+  bool prevGradMode() const {
+    return std::get<GradInterpreterMeta>(base_->meta()).prevGradMode_;
+  }
+  Tensor lift(const Tensor& tensor) const;
+ private:
+  const Interpreter* base_;
+};
+
+struct TORCH_API JvpInterpreterPtr {
+  explicit JvpInterpreterPtr(const Interpreter* base): base_(base) { TORCH_INTERNAL_ASSERT(base->key() == TransformType::Jvp); }
+  TransformType key() const { return base_->key(); }
+  int64_t level() const { return base_->level(); }
+  void processImpl(const c10::OperatorHandle& op, torch::jit::Stack* stack);
+  void sendToNextInterpreterImpl(const c10::OperatorHandle& op, torch::jit::Stack* stack, bool grad_special_case);
+  bool prevFwdGradMode() const {
+    return std::get<JvpInterpreterMeta>(base_->meta()).prevFwdGradMode_;
+  }
+  Tensor lift(const Tensor& tensor) const;
+ private:
+  const Interpreter* base_;
+};
+
+} // namespace at::functorch

pythonProject/.venv/Lib/site-packages/torch/include/ATen/functorch/BatchRulesHelper.h

@@ -0,0 +1,475 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+#pragma once
+
+#include <c10/util/TypeList.h>
+
+#include <ATen/ATen.h>
+#include <ATen/Operators.h>
+
+#include <ATen/functorch/DynamicLayer.h>
+#include <ATen/functorch/TensorWrapper.h>
+#include <ATen/functorch/BatchingMetaprogramming.h>
+#include <ATen/functorch/LegacyVmapTransforms.h>
+#include <ATen/functorch/BatchedFallback.h>
+#include <ATen/functorch/PlumbingHelper.h>
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <ATen/VmapGeneratedPlumbing.h>
+
+#include <utility>
+
+// This file contains helper functions for batching rules.
+
+namespace at::functorch {
+
+TORCH_API Tensor reshape_dim_into(int64_t src, int64_t dst, const Tensor& x);
+TORCH_API Tensor reshape_dim_outof(int64_t src, int64_t size1, const Tensor& x);
+
+TORCH_API Tensor reshape_dim_outof_symint(int64_t src, const c10::SymInt& size1, const Tensor& x);
+
+Tensor moveBatchDimToFront(const Tensor& tensor, optional<int64_t> maybe_batch_dim);
+int64_t rankWithoutBatchDim(const Tensor& tensor, optional<int64_t> maybe_batch_dim);
+int64_t numelWithoutBatchDim(const Tensor& tensor, optional<int64_t> maybe_batch_dim);
+optional<int64_t> valIfNonempty(optional<int64_t> maybe_empty, int64_t new_val);
+int64_t getPhysicalDim(const Tensor& tensor, bool has_batch_dim, int64_t logical_dim);
+VmapDimVector getPhysicalDims(const Tensor& tensor, bool has_batch_dim, IntArrayRef logical_dims);
+
+void vmapIncompatibleInplaceError(const char* schema_name);
+
+Tensor maybePadToLogicalRank(const Tensor& tensor, optional<int64_t> has_bdim, int64_t logical_rank);
+
+void check_randomness(RandomnessType randomness);
+void check_randomness(RandomnessType randomness, bool any_tensor_bdim);
+
+inline Tensor ensure_has_bdim(const Tensor& tensor, bool has_bdim, c10::SymInt batch_size) {
+  if (has_bdim) {
+    return tensor;
+  }
+  const auto sizes = tensor.sym_sizes();
+  SymDimVector expanded_shape;
+  expanded_shape.reserve(sizes.size());
+  expanded_shape.emplace_back(std::move(batch_size));
+  expanded_shape.insert(expanded_shape.end(), sizes.begin(), sizes.end());
+  return tensor.expand_symint(expanded_shape);
+}
+
+#define VMAP_SUPPORT(op, batch_rule) \
+  m.impl(#op, op ## _generated_plumbing<decltype(&batch_rule), &batch_rule>);
+
+#define VMAP_SUPPORT2(op, overload, batch_rule) \
+  m.impl(#op "." #overload, op ## _ ## overload ## _generated_plumbing<decltype(&batch_rule), &batch_rule>);
+
+#define OP_DECOMPOSE(op)  m.impl(#op, static_cast<decltype(&ATEN_FN(op))>(native::op));
+#define OP_DECOMPOSE2(op, overload)  m.impl(#op"."#overload, static_cast<decltype(&ATEN_FN2(op, overload))>(native::op));
+
+// DO NOT USE ME DIRECTLY! Use BASIC_UNARY_BATCH_RULE to save yourself some pain
+template <typename A, A a, typename C>
+struct BasicUnaryBatchRuleHelper;
+
+template <typename F, F Func, typename A, typename... T>
+struct BasicUnaryBatchRuleHelper<F, Func, c10::guts::typelist::typelist<A, T...>> {
+  static std::tuple<Tensor,optional<int64_t>> apply(
+      const Tensor& tensor,
+      optional<int64_t> batch_dim,
+      T... extra_args) {
+    return std::make_tuple(Func(tensor, std::forward<T>(extra_args)...), batch_dim);
+  }
+};
+
+// USAGE: BASIC_UNARY_BATCH_RULE(at::sin)
+// INCORRECT USAGE: BASIC_UNARY_BATCH_RULE(&at::sin)
+// It is important that this macro is not passed a function pointer!!
+#define BASIC_UNARY_BATCH_RULE(fn) SINGLE_ARG(\
+    BasicUnaryBatchRuleHelper<\
+      decltype(&fn),\
+      &fn,\
+      c10::guts::function_traits<decltype(fn)>::parameter_types>::apply)
+
+#define UNARY_POINTWISE(op) \
+  VMAP_SUPPORT(op, BASIC_UNARY_BATCH_RULE(ATEN_FN(op)));
+
+template <typename A, A a, typename C>
+struct VariadicBdimsBatchRuleHelper;
+
+template <typename F, F Func, typename A, typename... T>
+struct VariadicBdimsBatchRuleHelper<F, Func, c10::guts::typelist::typelist<A, T...>> {
+  static std::tuple<Tensor,optional<int64_t>> apply(
+      const Tensor& tensor,
+      optional<int64_t> batch_dim,
+      T... extra_args) {
+    auto tensor_ = moveBatchDimToFront(tensor, batch_dim);
+    return std::make_tuple(Func(tensor_, std::forward<T>(extra_args)...), 0);
+  }
+};
+
+// USAGE: VARIADIC_BDIMS_BATCH_RULE(at::cholesky_inverse)
+// INCORRECT USAGE: VARIADIC_BDIMS_BATCH_RULE(&at::cholesky_inverse)
+// It is important that this macro is not passed a function pointer!!
+#define VARIADIC_BDIMS_BATCH_RULE(fn) SINGLE_ARG(\
+    VariadicBdimsBatchRuleHelper<\
+      decltype(&fn),\
+      &fn,\
+      c10::guts::function_traits<decltype(fn)>::parameter_types>::apply)
+
+#define VARIADIC_BDIMS(op) \
+  VMAP_SUPPORT(op, VARIADIC_BDIMS_BATCH_RULE(ATEN_FN(op)));
+
+#define VARIADIC_BDIMS2(op, overload) \
+  VMAP_SUPPORT2(op, overload, VARIADIC_BDIMS_BATCH_RULE(ATEN_FN2(op, overload)));
+
+template<class F, F Func>
+void boxed_tensor_inputs_batch_rule(const c10::OperatorHandle& op, torch::jit::Stack* stack) {
+  const auto& schema = op.schema();
+  const auto num_returns = schema.returns().size();
+  const auto num_arguments = schema.arguments().size();
+
+  c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
+  auto maybe_layer = maybeCurrentDynamicLayer();
+  vmap_check_escaped(maybe_layer, "boxed_tensor_inputs_batch_rule");
+
+  int64_t cur_level = maybe_layer->layerId();
+
+  auto orig_arguments = torch::jit::last(*stack, num_arguments);
+  if (std::none_of(orig_arguments.begin(), orig_arguments.end(), ivalueParticipatesInCurrentLevel)) {
+    op.callBoxed(stack);
+    return;
+  }
+
+  auto arguments = torch::jit::pop(*stack, num_arguments);
+  std::vector<std::pair<Tensor, optional<int64_t>>> tensor_inputs;
+  std::vector<int64_t> tensor_pos;
+  for (const auto idx : c10::irange(0, num_arguments)) {
+    const auto& ivalue = arguments[idx];
+    if (ivalue.isTensor()) {
+      auto [tensor_value, tensor_bdim] = unwrapTensorAtLevel(ivalue.toTensor(), cur_level);
+      tensor_inputs.emplace_back(tensor_value, tensor_bdim);
+      tensor_pos.push_back(static_cast<int64_t>(idx));
+    }
+  }
+  Func(tensor_inputs);
+
+  size_t tensor_idx = 0;
+  TORCH_INTERNAL_ASSERT(!tensor_pos.empty());
+  for (const auto arg_idx : c10::irange(0, num_arguments)) {
+    if (tensor_idx >= tensor_pos.size() || (int64_t)arg_idx != tensor_pos[tensor_idx]) {
+      torch::jit::push(stack, arguments[arg_idx]);
+    } else {
+      TORCH_INTERNAL_ASSERT(tensor_idx < tensor_inputs.size());
+      torch::jit::push(stack, tensor_inputs[tensor_idx].first);
+      tensor_idx++;
+    }
+  }
+
+  op.callBoxed(stack);
+  const auto returns = torch::jit::pop(*stack, num_returns);
+  for (const auto& ret : returns) {
+    if (ret.isTensor()) {
+      torch::jit::push(stack, makeBatched(ret.toTensor(), 0, cur_level));
+    } else {
+      TORCH_INTERNAL_ASSERT(false, "This boxed batching rule does not currently support ops that return non-tensor values");
+    }
+  }
+}
+
+inline void handle_pointwise_ops(std::vector<std::pair<Tensor, optional<int64_t>>> &tensor_inputs) {
+  int64_t out_logical_rank = 0;
+  for (auto& tensor_input : tensor_inputs) {
+    int64_t cur_logical_rank = rankWithoutBatchDim(tensor_input.first, tensor_input.second);
+    out_logical_rank = std::max(out_logical_rank, cur_logical_rank);
+  }
+  for (auto& tensor_input: tensor_inputs) {
+    tensor_input.first = moveBatchDimToFront(tensor_input.first, tensor_input.second);
+    tensor_input.first = maybePadToLogicalRank(tensor_input.first, tensor_input.second, out_logical_rank);
+  }
+}
+
+#define POINTWISE_BOXED(op) \
+  m.impl(#op, torch::CppFunction::makeFromBoxedFunction<boxed_tensor_inputs_batch_rule<decltype(&handle_pointwise_ops), &handle_pointwise_ops>>());
+
+#define POINTWISE_BOXED2(op, overload) \
+  m.impl(#op "." #overload, torch::CppFunction::makeFromBoxedFunction<boxed_tensor_inputs_batch_rule<decltype(&handle_pointwise_ops), &handle_pointwise_ops>>());
+
+inline void handle_variadic_bdims(std::vector<std::pair<Tensor, optional<int64_t>>> &tensor_inputs) {
+  for (auto & tensor_input : tensor_inputs) {
+    tensor_input.first = moveBatchDimToFront(tensor_input.first, tensor_input.second);
+  }
+}
+
+#define VARIADIC_BDIMS_BOXED(op) \
+  m.impl(#op, torch::CppFunction::makeFromBoxedFunction<boxed_tensor_inputs_batch_rule<decltype(&handle_variadic_bdims), &handle_variadic_bdims>>());
+
+using UnpackedBatchedTensor = std::tuple<Tensor,optional<int64_t>>;
+
+inline void find_and_unpack_tensors(
+    const torch::jit::Stack* stack,
+    int64_t num_args,
+    int64_t cur_level,
+    SmallVector<UnpackedBatchedTensor, 5>* tensors,
+    SmallVector<int64_t, 5>* tensors_pos,
+    int64_t* batch_size) {
+
+  int64_t computed_batch_size = -1;
+  int64_t args_begin = static_cast<int64_t>(stack->size()) - num_args;
+
+  for (const auto idx : c10::irange(0, num_args)) {
+    const auto& ivalue = (*stack)[args_begin + idx];
+    if (!ivalue.isTensor()) {
+      continue;
+    }
+    auto unpacked = unwrapTensorAtLevel(ivalue.toTensor(), cur_level);
+    const auto& tensor_value = std::get<0>(unpacked);
+    const auto tensor_bdim = std::get<1>(unpacked);
+    if (tensor_bdim.has_value()) {
+      auto candidate_batch_size = tensor_value.size(*tensor_bdim);
+      if (computed_batch_size == -1) {
+        computed_batch_size = candidate_batch_size;
+      }
+      TORCH_INTERNAL_ASSERT(candidate_batch_size == computed_batch_size);
+    }
+
+    tensors->push_back(std::move(unpacked));
+    tensors_pos->push_back(idx);
+  }
+  TORCH_INTERNAL_ASSERT(computed_batch_size > -1);
+  *batch_size = computed_batch_size;
+}
+
+inline void boxed_existing_bdim_all_batch_rule(
+    const c10::OperatorHandle& op, torch::jit::Stack* stack) {
+  const auto& schema = op.schema();
+  const auto num_returns = schema.returns().size();
+  const auto num_arguments = static_cast<int64_t>(schema.arguments().size());
+
+  c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
+  auto maybe_layer = maybeCurrentDynamicLayer();
+  vmap_check_escaped(maybe_layer, "boxed_existing_bdim_all_batch_rule");
+  int64_t cur_level = maybe_layer->layerId();
+
+  const auto arguments = torch::jit::last(stack, num_arguments);
+  if (std::none_of(arguments.begin(), arguments.end(), ivalueParticipatesInCurrentLevel)) {
+    op.callBoxed(stack);
+    return;
+  }
+
+  int64_t args_begin = static_cast<int64_t>(stack->size()) - num_arguments;
+  SmallVector<UnpackedBatchedTensor, 5> tensor_inputs;
+  SmallVector<int64_t, 5> tensor_pos;
+  int64_t batch_size = 0;
+
+  find_and_unpack_tensors(
+      stack, num_arguments, cur_level,
+      &tensor_inputs, &tensor_pos, &batch_size);
+
+  // for each tensor, ensure it has a bdim and reshape it.
+  for (const auto tensor_idx : c10::irange(0, tensor_inputs.size())) {
+    const auto& value = std::get<0>(tensor_inputs[tensor_idx]);
+    auto bdim = std::get<1>(tensor_inputs[tensor_idx]);
+    auto value_ = ensure_has_bdim(value, bdim.has_value(), batch_size);
+    if (!bdim.has_value()) {
+      bdim = 0;
+    }
+    (*stack)[args_begin + tensor_pos[tensor_idx]] = reshape_dim_into(*bdim, 0, value_);
+  }
+
+  op.callBoxed(stack);
+
+  for (const auto idx : c10::irange(args_begin, args_begin + num_returns)) {
+    const auto& ret = (*stack)[idx];
+    TORCH_INTERNAL_ASSERT(ret.isTensor(),
+        "This boxed batching rule does not currently support ops that return non-tensor values");
+    (*stack)[idx] = makeBatched(reshape_dim_outof(0, batch_size, ret.toTensor()), 0, cur_level);
+  }
+}
+
+// Use when all tensors arguments accept one (normal) batch dim.
+// This batching rule expands the batch dim on all Tensors, reshapes it into
+// dim 0, calls the op, and then reshapes the batch dim out of dim 0.
+// This is not the most efficient thing; if there are alternatives, plese try
+// to use them. Use this only as a last resort.
+#define EXISTING_BDIM_ALL_BOXED(op) \
+  m.impl(#op, torch::CppFunction::makeFromBoxedFunction<boxed_existing_bdim_all_batch_rule>());
+
+template <int64_t feature_rank, int64_t contig_tensor_index=-1>
+inline void boxed_all_tensors_have_optional_bdim(
+    const c10::OperatorHandle& op, torch::jit::Stack* stack) {
+  const auto& schema = op.schema();
+  const auto num_returns = schema.returns().size();
+  const auto num_arguments = schema.arguments().size();
+
+  c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
+  auto maybe_layer = maybeCurrentDynamicLayer();
+  vmap_check_escaped(maybe_layer, "boxed_all_tensors_have_optional_bdim");
+  int64_t cur_level = maybe_layer->layerId();
+
+  const auto arguments = torch::jit::last(stack, num_arguments);
+  if (std::none_of(arguments.begin(), arguments.end(), ivalueParticipatesInCurrentLevel)) {
+    op.callBoxed(stack);
+    return;
+  }
+
+  int64_t args_begin = static_cast<int64_t>(stack->size() - num_arguments);
+  SmallVector<UnpackedBatchedTensor, 5> tensor_inputs;
+  SmallVector<int64_t, 5> tensor_pos;
+  int64_t batch_size = 0;
+
+  find_and_unpack_tensors(
+      stack, static_cast<int64_t>(num_arguments), cur_level,
+      &tensor_inputs, &tensor_pos, &batch_size);
+
+  optional<bool> is_no_batch_dim_case;
+
+  for (const auto tensor_idx : c10::irange(0, tensor_inputs.size())) {
+    const auto& value = std::get<0>(tensor_inputs[tensor_idx]);
+    auto bdim = std::get<1>(tensor_inputs[tensor_idx]);
+    const auto logical_rank = rankWithoutBatchDim(value, bdim);
+
+    if (!is_no_batch_dim_case.has_value()) {
+      is_no_batch_dim_case = (logical_rank == feature_rank);
+    }
+    auto value_ = ensure_has_bdim(value, bdim.has_value(), batch_size);
+    if (!bdim.has_value()) {
+      bdim = 0;
+    }
+    if (*is_no_batch_dim_case) {
+      TORCH_INTERNAL_ASSERT(logical_rank == feature_rank);
+      value_ = moveBatchDimToFront(value_, bdim);
+      if (tensor_idx == contig_tensor_index) {
+        value_ = value_.contiguous();
+      }
+      (*stack)[args_begin + tensor_pos[tensor_idx]] = std::move(value_);
+      continue;
+    }
+    TORCH_INTERNAL_ASSERT(logical_rank == feature_rank + 1);
+    value_ = reshape_dim_into(*bdim, 0, value_);
+    if (tensor_idx == contig_tensor_index) {
+      value_ = value_.contiguous();
+    }
+    (*stack)[args_begin + tensor_pos[tensor_idx]] = std::move(value_);
+  }
+
+  op.callBoxed(stack);
+
+  for (const auto idx : c10::irange(args_begin, args_begin + num_returns)) {
+    const auto& ret = (*stack)[idx];
+    TORCH_INTERNAL_ASSERT(ret.isTensor(),
+        "This boxed batching rule does not currently support ops that return non-tensor values");
+    if (*is_no_batch_dim_case) {
+      (*stack)[idx] = makeBatched(ret.toTensor(), 0, cur_level);
+    } else {
+      (*stack)[idx] = makeBatched(reshape_dim_outof(0, batch_size, ret.toTensor()), 0, cur_level);
+    }
+  }
+}
+
+// Useful for many NN operators.
+// The operator must satisfy the following:
+// - All arguments must accept an optional batch dim.
+// - All arguments must be the same rank
+#define ALL_TENSORS_HAVE_OPTIONAL_BDIM_BOXED(feature_rank, op) \
+  m.impl(#op, torch::CppFunction::makeFromBoxedFunction<boxed_all_tensors_have_optional_bdim<feature_rank>>());
+
+#define ALL_TENSORS_HAVE_OPTIONAL_BDIM_BOXED_CONTIG1(feature_rank, op, contig_tensor_index) \
+  m.impl(#op, \
+         torch::CppFunction::makeFromBoxedFunction<\
+             boxed_all_tensors_have_optional_bdim<\
+                 feature_rank, \
+                 contig_tensor_index>\
+             >());
+
+template <typename A, A a, typename C>
+struct ExistingBdimBatchRuleHelper;
+
+template <typename F, F Func, typename A, typename... T>
+struct ExistingBdimBatchRuleHelper<F, Func, c10::guts::typelist::typelist<A, T...>> {
+  static std::tuple<Tensor,optional<int64_t>> apply(
+      const Tensor& self,
+      optional<int64_t> self_bdim,
+      T... extra_args) {
+    auto self_ = reshape_dim_into(*self_bdim, 0, self);
+    auto out = Func(self_, std::forward<T>(extra_args)...);
+    return std::make_tuple(reshape_dim_outof_symint(0, self.sym_sizes()[*self_bdim], out), 0);
+  }
+};
+
+// USAGE: EXISTING_BDIM_BATCH_RULE(at::cholesky_inverse)
+// INCORRECT USAGE: EXISTING_BDIM_BATCH_RULE(&at::cholesky_inverse)
+// It is important that this macro is not passed a function pointer!!
+#define EXISTING_BDIM_BATCH_RULE(fn) SINGLE_ARG(\
+    ExistingBdimBatchRuleHelper<\
+      decltype(&fn),\
+      &fn,\
+      c10::guts::function_traits<decltype(fn)>::parameter_types>::apply)
+
+
+#define EXISTING_BDIM(op) \
+  VMAP_SUPPORT(op, EXISTING_BDIM_BATCH_RULE(ATEN_FN(op)));
+
+#define EXISTING_BDIM2(op, overload) \
+  VMAP_SUPPORT2(op, overload, EXISTING_BDIM_BATCH_RULE(ATEN_FN2(op, overload)));
+
+#define INVOKE(object,ptrToMember)  ((object).*(ptrToMember))
+
+
+template <typename F, F Method, typename... ExtraArgs>
+Tensor& unary_inplace_batch_rule(Tensor& self, optional<int64_t>, ExtraArgs... extra_args) {
+  INVOKE(self, Method)(std::forward<ExtraArgs>(extra_args)...);
+  return self;
+}
+
+inline int64_t get_bdim_size4(
+    const Tensor& a_value, optional<int64_t> a_bdim,
+    const Tensor& b_value, optional<int64_t> b_bdim,
+    const Tensor& c_value, optional<int64_t> c_bdim,
+    const Tensor& d_value, optional<int64_t> d_bdim) {
+  if (a_bdim)
+    return a_value.size(*a_bdim);
+  if (b_bdim)
+    return b_value.size(*b_bdim);
+  if (c_bdim)
+    return c_value.size(*c_bdim);
+  if (d_bdim)
+    return d_value.size(*d_bdim);
+  TORCH_INTERNAL_ASSERT(false);
+}
+
+inline int64_t get_bdim_size3(
+    const Tensor& a_value, optional<int64_t> a_bdim,
+    const Tensor& b_value, optional<int64_t> b_bdim,
+    const Tensor& c_value, optional<int64_t> c_bdim) {
+  if (a_bdim)
+    return a_value.size(*a_bdim);
+  if (b_bdim)
+    return b_value.size(*b_bdim);
+  if (c_bdim)
+    return c_value.size(*c_bdim);
+  TORCH_INTERNAL_ASSERT(false);
+}
+
+inline int64_t get_bdim_size2(
+    const Tensor& a_value, optional<int64_t> a_bdim,
+    const Tensor& b_value, optional<int64_t> b_bdim) {
+  if (a_bdim)
+    return a_value.size(*a_bdim);
+  if (b_bdim)
+    return b_value.size(*b_bdim);
+  TORCH_INTERNAL_ASSERT(false);
+}
+
+// [start, start + 1, ..., stop - 1]
+inline VmapDimVector range(int64_t start, int64_t stop) {
+  TORCH_INTERNAL_ASSERT(stop >= start);
+  VmapDimVector dims;
+  dims.reserve(stop - start);
+  for (int64_t i = start; i < stop; i++) {
+    dims.emplace_back(i);
+  }
+  return dims;
+}
+std::tuple<Tensor, Tensor> _binary_pointwise_helper(
+    const Tensor& tensor, optional<int64_t> tensor_batch_dim, const Tensor& other, optional<int64_t> other_batch_dim,
+    bool do_type_promotion=true);
+
+} // namespace at::functorch

pythonProject/.venv/Lib/site-packages/torch/include/ATen/functorch/BatchedFallback.h

@@ -0,0 +1,81 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+#include <ATen/ATen.h>
+#include <ATen/core/op_registration/op_registration.h>
+#include <torch/library.h>
+
+namespace at::functorch {
+
+// This file contains code for the vmap fallback (also known as the
+// BatchedTensor fallback or the Batched fallback). This code runs
+// when an operation doesn't have a batching rule implemented.
+
+// If an operator doesn't have a batching rule implemented then we fallback
+// to this implementation. The fallback doesn't work on out= variants or
+// view operations; that is, it works for out-of-place operations and
+// in-place non-view operations.
+//
+// For out-of-place operations, the fallback effectively takes all of the
+// BatchedTensors in `stack`, slices them, and runs `op` on all of the
+// corresponding slices to produce slices of the outputs. The output slices
+// then get `torch.stack`ed to create the
+// final returns.
+//
+// The performance of the fallback is not very good because it introduces an
+// extra copy from stacking the sliced outputs. Because of this, we prefer to
+// write batching rules for operators whenever possible.
+void batchedTensorForLoopFallback(const c10::OperatorHandle& op, torch::jit::Stack* stack);
+void batchedNestedTensorForLoopFallback(const c10::OperatorHandle& op, torch::jit::Stack* stack);
+
+void vmapErrorFallback(const c10::OperatorHandle& op, torch::jit::Stack* stack);
+
+// The vmap fallback emits a warning by default, but it may be disabled if
+// the user finds it to be too annoying.
+TORCH_API bool isVmapFallbackWarningEnabled();
+TORCH_API void setVmapFallbackWarningEnabled(bool enabled);
+
+// Used for testing. The vmap fallback is enabled by default. When it is disabled,
+// it raises an error.
+TORCH_API bool isVmapFallbackEnabled();
+TORCH_API void setVmapFallbackEnabled(bool enabled);
+
+template <typename A> A vector_to_result(const std::vector<IValue>& buffer) {
+  return buffer[0].to<A>();
+}
+template <typename A, typename B> std::tuple<A, B> vector_to_result(const std::vector<IValue>& buffer) {
+  return std::make_tuple(buffer[0].to<A>(), buffer[1].to<B>());
+}
+template <typename A, typename B, typename C> std::tuple<A, B, C> vector_to_result(const std::vector<IValue>& buffer) {
+  return std::make_tuple(buffer[0].to<A>(), buffer[1].to<B>(), buffer[2].to<B>());
+}
+
+// slow_fallback is a way to call the vmap fallback inside some boxed kernel.
+// There is probably some better way to metaprogram this.
+template <typename Ret>
+Ret slow_fallback(const c10::OperatorHandle& op, ArrayRef<IValue> args) {
+  std::vector<IValue> stack(args.begin(), args.end());
+  batchedTensorForLoopFallback(op, &stack);
+  return vector_to_result<Ret>(stack);
+}
+
+template <typename A, typename B>
+std::tuple<A, B> slow_fallback(const c10::OperatorHandle& op, ArrayRef<IValue> args) {
+  std::vector<IValue> stack(args.begin(), args.end());
+  batchedTensorForLoopFallback(op, &stack);
+  return vector_to_result<A, B>(stack);
+}
+
+template <typename A, typename B, typename C>
+std::tuple<A, B, C> slow_fallback(const c10::OperatorHandle& op, ArrayRef<IValue> args) {
+  std::vector<IValue> stack(args.begin(), args.end());
+  batchedTensorForLoopFallback(op, &stack);
+  return vector_to_result<A, B, C>(stack);
+}
+
+
+} // namespace at::functorch

pythonProject/.venv/Lib/site-packages/torch/include/ATen/functorch/BatchedTensorImpl.h

@@ -0,0 +1,169 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <bitset>
+
+#include <ATen/ArrayRef.h>
+#include <ATen/SmallVector.h>
+#include <ATen/Tensor.h>
+
+namespace at::functorch {
+
+using Tensor = at::Tensor;
+
+// We assume this in a few other places in the codebase,
+// but there isn't a centralized definition.
+constexpr int64_t kVmapMaxTensorDims = 64;
+
+// The valid vmap levels range from [0, 64). This effectively means that we
+// support a maximum of 64 nested vmaps.
+constexpr int64_t kVmapNumLevels = 64;
+
+// Store this number of elements of BatchDims on the stack. Most people will
+// probably use <= 5 nested vmaps, but adjust this number as necessary.
+constexpr int64_t kBatchDimsStackSize = 5;
+
+// A BatchedTensorImpl holds an underlying Tensor and a single batch dim
+// NB: We use the term "BatchedTensor" to mean a Tensor that is backed with a
+// BatchedTensorImpl.
+//
+// The batch dimensions are treated as being "private"; they are not user-visible.
+// For example, in the following Tensor,
+//    bt = BatchedTensorImpl(ones(2, 3, 5, 7), lvl=1, dim=0)
+// dimension 0 is batch dimension.
+//
+// bt.sizes() returns (5, 7); bt.sum(0) performs a reduction over the (public)
+// dim 0, which is equivalent to dim 3 in the underlying ones(2, 3, 5, 7) tensor.
+struct TORCH_API BatchedTensorImpl : public c10::TensorImpl {
+  explicit BatchedTensorImpl(at::DispatchKeySet key_set, Tensor value, int64_t dim, int64_t level);
+
+  // Returns batch dimension of this tensor
+  int64_t bdim() const { return bdim_; }
+
+  // Returns batch dimension of this tensor
+  int64_t level() const { return level_; }
+
+  // BatchedTensorImpl wraps a Tensor
+  const Tensor& value() const { return value_; }
+
+  // Given a public dimension index, return the dimension index in the underlying
+  // value() tensor.
+  // For example, if we have
+  //    bt = BatchedTensorImpl(ones(2, 3, 5, 7), lvl=1, dim=0)
+  // bt.actualDim(0) -> 1
+  // bt.actualDim(1) -> 2
+  // bt.actualDim(2) -> 3
+  // bt.actualDim(3) -> Error
+  int64_t actualDim(int64_t dim, bool wrap_dim = true) const;
+
+  IntArrayRef sizes_custom() const override;
+  SymIntArrayRef sym_sizes_custom() const override;
+  int64_t size_custom(int64_t d) const override;
+  c10::SymInt sym_size_custom(int64_t d) const override;
+  // We have to override this because we opted into CustomStrides
+  IntArrayRef strides_custom() const override;
+  SymIntArrayRef sym_strides_custom() const override;
+  // Override a bunch of methods inherited from TensorImpl to return error messages.
+  bool is_contiguous_custom(at::MemoryFormat memory_format=at::MemoryFormat::Contiguous) const override;
+  void set_size(int64_t dim, int64_t new_size) override;
+  void set_stride(int64_t dim, int64_t new_stride) override;
+  c10::intrusive_ptr<TensorImpl> shallow_copy_and_detach(
+    const c10::VariableVersion& version_counter,
+    bool allow_tensor_metadata_change) const override;
+  c10::intrusive_ptr<TensorImpl> shallow_copy_and_detach(
+      c10::VariableVersion&& version_counter,
+      bool allow_tensor_metadata_change) const override;
+  void shallow_copy_from(const c10::intrusive_ptr<TensorImpl>& impl) override;
+#ifdef DEBUG
+  bool has_storage() const override;
+#endif
+
+  void refreshTensorMetadata();
+
+  // Used in torchdim. torchdim uses non-lexical BatchedTensor; the way it
+  // accomplishes this is a hack where it is able to modify the levels of
+  // BatchedTensor to match the level of the current vmap transform.
+  void _unsafe_set_level(int64_t level) {
+    level_ = level;
+  }
+
+  // Used in batching rule for in-place view operations that can change
+  // the index of the bdim (think squeeze_, unsqueeze_)
+  void unsafe_set_bdim(int64_t bdim) {
+    // NB: you MUST call refreshTensorMetadata after doing this.
+    bdim_ = bdim;
+  }
+ private:
+  // see NOTE: [BatchedTensorImpl levels invariant]
+  void checkInvariants() const;
+  const char* tensorimpl_type_name() const override;
+
+  Tensor value_;
+
+  int64_t level_;
+  int64_t bdim_;
+};
+
+// NB: We use the term "BatchedTensor" to mean a Tensor that is backed with a
+// BatchedTensorImpl.
+inline bool isBatchedTensor(const Tensor& tensor) {
+  return tensor.unsafeGetTensorImpl()->key_set().has(DispatchKey::FuncTorchBatched) ||
+      tensor.unsafeGetTensorImpl()->key_set().has(DispatchKey::BatchedNestedTensor);
+}
+
+// It is unsafe to call this on a Tensor that is not backed by a
+// BatchedTensorImpl. Please use `maybeGetBatchedImpl` whenever possible.
+inline BatchedTensorImpl* unsafeGetBatchedImpl(const Tensor& tensor) {
+  return static_cast<BatchedTensorImpl*>(tensor.unsafeGetTensorImpl());
+}
+
+inline BatchedTensorImpl* maybeGetBatchedImpl(const Tensor& tensor) {
+  if (!isBatchedTensor(tensor)) {
+    return nullptr;
+  }
+  return unsafeGetBatchedImpl(tensor);
+}
+
+// Returns a bitset. If bit i is set, then that means dim i is a batchdim.
+inline std::bitset<kVmapMaxTensorDims> createBatchDimBitset(int64_t dim) {
+  std::bitset<kVmapMaxTensorDims> is_bdim;
+  is_bdim.set(dim);
+  return is_bdim;
+}
+
+// Creates a bitset for the given level
+inline std::bitset<kVmapNumLevels> createVmapLevelsBitset(int64_t level) {
+  std::bitset<kVmapNumLevels> result;
+  result.set(level);
+  return result;
+}
+
+// Use this to construct a BatchedTensor from a regular Tensor
+TORCH_API Tensor makeBatched(const Tensor& tensor, int64_t dim, int64_t level);
+
+// Adds a batch dim to `tensor`, returning a BatchedTensor
+TORCH_API Tensor addBatchDim(const Tensor& tensor, int64_t dim, int64_t level);
+
+// Certain dispatch keys must be propagated to the BatchedTensor (or, in general,
+// any wrapper Tensor subclasses). This is because there are methods on Tensor
+// that skip dispatch and check for the presence of a dispatch key (e.g. is_cpu()).
+// TODO: should probably contain more (or all?) backend keys
+constexpr DispatchKeySet kKeysToPropagateToWrapper({
+  DispatchKey::Negative,
+  DispatchKey::Conjugate,
+  DispatchKey::XLA,
+  DispatchKey::CUDA,
+  DispatchKey::CPU,
+});
+
+inline DispatchKeySet getKeysToPropagateToWrapper(const Tensor& tensor, DispatchKeySet to_propagate=kKeysToPropagateToWrapper) {
+  auto key_set = tensor.unsafeGetTensorImpl()->key_set();
+  return key_set & kKeysToPropagateToWrapper;
+}
+
+} // namespace at::functorch

pythonProject/.venv/Lib/site-packages/torch/include/ATen/functorch/BatchingMetaprogramming.h

@@ -0,0 +1,126 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+#include <ATen/Tensor.h>
+#include <ATen/VmapGeneratedPlumbing.h>
+
+// This file contains template metaprogramming things that are used for our
+// batching rules.
+//
+// See NOTE: [vmap plumbing] for more details on why this is necessary.
+// The plumbing has a bunch of metaprogramming hacks for determining the signature
+// of a batching rule from the signature of the operator, many of which use the
+// helper functions in this file.
+
+namespace at::functorch {
+
+// Metaprogramming things
+template <class... Items> using typelist = c10::guts::typelist::typelist<Items...>;
+template <class TypeList> using head_t = c10::guts::typelist::head_t<TypeList>;
+template <class TL1, class TL2> using concat_t = c10::guts::typelist::concat_t<TL1, TL2>;
+template <typename T> class debug_t;
+
+// tail operation
+template<class TypeList>
+struct tail final {
+    static_assert(c10::guts::false_t<TypeList>::value,
+                  "In typelist::tail<T>, the T argument must be typelist<...>.");
+};
+template<class Head, class... Tail>
+struct tail<typelist<Head, Tail...>> final {
+  using type = typelist<Tail...>;
+};
+template<class TypeList> using tail_t = typename tail<TypeList>::type;
+
+template <class First, class Second, class Next, class Tail>
+struct IfFirstIsTensorAndSecondisBatchDimThenTailElseNext {
+  using type = Next;
+};
+template <class Next, class Tail>
+struct IfFirstIsTensorAndSecondisBatchDimThenTailElseNext<Tensor, optional<int64_t>, Next, Tail> {
+  using type = Tail;
+};
+template <class Next, class Tail>
+struct IfFirstIsTensorAndSecondisBatchDimThenTailElseNext<const Tensor&, optional<int64_t>, Next, Tail> {
+  using type = Tail;
+};
+template <class Next, class Tail>
+struct IfFirstIsTensorAndSecondisBatchDimThenTailElseNext<Tensor&, optional<int64_t>, Next, Tail> {
+  using type = Tail;
+};
+template <class Next, class Tail>
+struct IfFirstIsTensorAndSecondisBatchDimThenTailElseNext<optional<Tensor>, optional<int64_t>, Next, Tail> {
+  using type = Tail;
+};
+template <class Next, class Tail>
+struct IfFirstIsTensorAndSecondisBatchDimThenTailElseNext<const optional<Tensor>&, optional<int64_t>, Next, Tail> {
+  using type = Tail;
+};
+template <class Next, class Tail>
+struct IfFirstIsTensorAndSecondisBatchDimThenTailElseNext<optional<Tensor>&, optional<int64_t>, Next, Tail> {
+  using type = Tail;
+};
+template <class Next, class Tail>
+struct IfFirstIsTensorAndSecondisBatchDimThenTailElseNext<std::vector<Tensor>, optional<int64_t>, Next, Tail> {
+  using type = Tail;
+};
+template <class TypeList> struct RemoveBatchDimAfterTensor {
+  using first = head_t<TypeList>;
+  using next = tail_t<TypeList>;
+  using second = head_t<next>;
+  using tail = tail_t<next>;
+
+  using type = concat_t<
+    typelist<first>,
+    typename RemoveBatchDimAfterTensor<
+      typename IfFirstIsTensorAndSecondisBatchDimThenTailElseNext<first, second, next, tail>::type
+    >::type
+  >;
+};
+template <class Type> struct RemoveBatchDimAfterTensor<typelist<Type>> {
+  using type = typelist<Type>;
+};
+template <> struct RemoveBatchDimAfterTensor<typelist<>> {
+  using type = typelist<>;
+};
+template<class TypeList> using remove_batch_dim_after_tensor_t = typename RemoveBatchDimAfterTensor<TypeList>::type;
+
+template <typename T> struct UnpackSingleItemTuple {
+  using type = T;
+};
+template <typename T> struct UnpackSingleItemTuple<std::tuple<T>> {
+  using type = T;
+};
+template <typename T> using unpack_single_item_tuple_t = typename UnpackSingleItemTuple<T>::type;
+
+template <typename Return, typename TupleArgs> struct BuildFunctionHelper;
+template <typename Return, typename... Args> struct BuildFunctionHelper<Return, std::tuple<Args...>> {
+  using type = Return(Args...);
+};
+template <typename Return, typename TL>
+struct BuildFunction {
+  using type = typename BuildFunctionHelper<Return, c10::guts::typelist::to_tuple_t<TL>>::type;
+};
+template <typename Return, typename TL> using build_function_t = typename BuildFunction<Return, TL>::type;
+
+
+template <typename batch_rule_t> struct ToOperatorType {
+  using batch_rule_return_type = typename c10::guts::function_traits<batch_rule_t>::return_type;
+  using batch_rule_parameter_types = typename c10::guts::function_traits<batch_rule_t>::parameter_types;
+
+  using operator_parameter_types = remove_batch_dim_after_tensor_t<batch_rule_parameter_types>;
+  using operator_return_type =
+    unpack_single_item_tuple_t<
+      c10::guts::typelist::to_tuple_t<
+        remove_batch_dim_after_tensor_t<
+          c10::guts::typelist::from_tuple_t<batch_rule_return_type>>>>;
+
+  using type = build_function_t<operator_return_type, operator_parameter_types>;
+};
+template <typename batch_rule_t> using to_operator_t = typename ToOperatorType<batch_rule_t>::type;
+
+} // namespace at::functorch

pythonProject/.venv/Lib/site-packages/torch/include/ATen/functorch/DynamicLayer.h

@@ -0,0 +1,124 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+#include <ATen/functorch/Macros.h>
+#include <c10/core/DispatchKey.h>
+#include <ATen/core/function_schema.h>
+#include <c10/util/Optional.h>
+#include <c10/core/impl/LocalDispatchKeySet.h>
+#include <ATen/functorch/Interpreter.h>
+#include <ATen/functorch/VmapInterpreter.h>
+#include <ATen/functorch/ADInterpreters.h>
+#include <ATen/functorch/FunctionalizeInterpreter.h>
+
+// Forward declared
+namespace c10 { struct AutogradMetaInterface; }
+
+namespace at::functorch  {
+
+// This file contains the implementation of functorch's interpreter stack.
+// See NOTE: [functorch interpreter stack] first before reading on.
+//
+// NB: the functorch interpreter stack is also referred to as:
+// - the "dynamic layer stack" -- an older name for "interpreter" was
+//   "dynamic layer".
+// - the "functorch mode stack". You can think of each functorch transform as a
+//   "mode" (in the same sense as torch_dispatch mode or torch_function mode),
+//   and functorch being an implementation of a "mode stack" where the modes
+//   may be arbitrary composed.
+
+// DynamicLayer is basically the same thing as an Interpreter.
+// It represents a functorch transform and it holds an Interpreter,
+// which contains metadata related to the transform and instructions on
+// how to perform the transform.
+//
+// TODO: we can excise DynamicLayer in favor of Interpreter,
+// But I am going to leave it for now as a compatiblity shim to avoid
+// needing to refactor a lot of callsites...
+struct TORCH_API DynamicLayer {
+  explicit DynamicLayer(
+      TransformType transform_type,
+      int64_t layerId,
+      optional<c10::SymInt> batchSize = nullopt,
+      optional<RandomnessType> randomness = nullopt,
+      optional<bool> prev_grad_mode = nullopt,
+      optional<bool> pre_fwd_grad_mode = nullopt,
+      optional<bool> functionalize_add_back_views = nullopt);
+
+  TransformType key() const;
+  int64_t layerId() const;
+
+  const Interpreter& interpreter() const { return interpreter_; }
+  Interpreter& interpreter() { return interpreter_; }
+
+  // Only valid for vmap
+  c10::SymInt batchSize() const;
+  RandomnessType randomness() const;
+
+ private:
+  Interpreter interpreter_;
+};
+
+TORCH_API int64_t initAndPushDynamicLayer(
+    TransformType transform_type,
+    optional<c10::SymInt> batch_size = nullopt,
+    optional<RandomnessType> randomness = nullopt,
+    optional<bool> prev_grad_mode = nullopt,
+    optional<bool> prev_fwd_grad_mode = nullopt,
+    optional<bool> functionalize_add_back_views = nullopt);
+TORCH_API DynamicLayer popDynamicLayerAndDeleteMetadata();
+TORCH_API std::optional<DynamicLayer> maybeCurrentDynamicLayer();
+TORCH_API const std::vector<DynamicLayer>& getDynamicLayerStack();
+TORCH_API void setDynamicLayerStack(const std::vector<DynamicLayer>& stack);
+TORCH_API void setDynamicLayerFrontBackKeysIncluded(bool included);
+
+// NOTE: [Life handles and lexically scoped transforms]
+// functorch transforms are lexically scoped.
+// Given a level, we store a "life handle" that is a boolean that tells us if the
+// transform with that level is active or not.
+//
+// functorch's TensorWrapper (for grad transforms) stores a life handle.
+// If a TensorWrapper escapes from the scope of the transform, then somehow
+// it must know it escaped; it can tell by querying the life handle.
+TORCH_API const std::shared_ptr<bool>& getLifeHandleForLevel(int64_t level);
+
+// Returns if an operator is in-place. An operator is inplace if:
+// 1. The first argument is a Tensor and it is being written to
+// 2. The first argument is being returned
+// 3. No other arguments are aliased
+// Here is an example of an in-place operator:
+// add_(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> Tensor(a!)
+TORCH_API bool isInplaceOp(const c10::FunctionSchema& schema);
+
+// Given the indices of unwrapped inputs and the schema, this returns the indices of any outputs that should remain unwrapped
+TORCH_API std::optional<size_t> findAliasedOutput(const FunctionSchema& schema, const int64_t immutable_input);
+
+TORCH_API Tensor unwrapIfDead(const Tensor& tensor);
+TORCH_API bool isDeadTensorWrapper(const Tensor& tensor);
+
+// Pretty printers
+TORCH_API std::ostream& operator<<(std::ostream& os, const DynamicLayer& layer);
+TORCH_API std::ostream& operator<<(std::ostream& os, const std::vector<DynamicLayer>& dynamicLayerStack);
+
+// While a functorch transform is active, torch.autograd.function._SingleLevelFunction
+// is disabled by default. The following two APIs are APIs for enabling
+// it. These are not user-facing APIs. We can delete this in the future, but
+// it is useful for debugging when something goes wrong with the
+// autograd.Function <> functorch interaction, which uses _SingleLevelFunction,
+// because it leads to loud errors if something is incorrect.
+TORCH_API void setSingleLevelAutogradFunctionAllowed(bool allowed);
+TORCH_API bool getSingleLevelAutogradFunctionAllowed();
+
+// While a functorch grad transform is active, Tensor.requires_grad_() gets
+// disabled. These two functions are the mechanism to controlling that.
+TORCH_API void setInplaceRequiresGradAllowed(bool allowed);
+TORCH_API bool getInplaceRequiresGradAllowed();
+
+TORCH_API DynamicLayer popDynamicLayer();
+TORCH_API int64_t pushDynamicLayer(DynamicLayer&& layer);
+
+} // namespace at::functorch

pythonProject/.venv/Lib/site-packages/torch/include/ATen/functorch/FunctionalizeInterpreter.h

@@ -0,0 +1,22 @@
+#pragma once
+#include <ATen/functorch/Interpreter.h>
+
+namespace at::functorch {
+
+// This is the interpreter that handles the functionalize() transform.
+// See NOTE: [functorch interpreter stack] for more details.
+
+struct FunctionalizeInterpreterPtr {
+  explicit FunctionalizeInterpreterPtr(const Interpreter* base): base_(base) { TORCH_INTERNAL_ASSERT(base->key() == TransformType::Functionalize); }
+  TransformType key() const { return base_->key(); }
+  int64_t level() const { return base_->level(); }
+  void processImpl(const c10::OperatorHandle& op, torch::jit::Stack* stack);
+  void sendToNextInterpreterImpl(const c10::OperatorHandle& op, torch::jit::Stack* stack, bool grad_special_case);
+  bool functionalizeAddBackViews() const {
+    return std::get<FunctionalizeInterpreterMeta>(base_->meta()).functionalizeAddBackViews_;
+  }
+ private:
+  const Interpreter* base_;
+};
+
+} // namespace at::functorch

pythonProject/.venv/Lib/site-packages/torch/include/ATen/functorch/Interpreter.h

@@ -0,0 +1,209 @@
+#pragma once
+
+#include <ATen/functorch/Macros.h>
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <c10/core/impl/LocalDispatchKeySet.h>
+#include <c10/util/Optional.h>
+#include <bitset>
+#include <utility>
+#include <variant>
+
+namespace at::functorch {
+
+// NOTE: [functorch interpreter stack]
+//
+// functorch's dispatching system uses a stack of interpreters.
+// Historically we've referred to this as the "DynamicLayerStack".
+//
+// An interpreter is something that reads in the code it is passed
+// and then executes it. We have a different interpreter per-transform:
+// the "VmapInterpreter" is responsible for reading in operators (like aten::mv)
+// and executing the batched version of it (the batching rule for aten::mv).
+//
+// Concretely, each interpreter is responsible for two things:
+//
+// 1) process(ophandle, stack)
+// Given an operator handle and a stack of arguments, the interpreter is
+// responsible for figuring out how to execute the operation under the semantics
+// of the interpreter. For e.g. VmapInterpreter, this is figuring out how to call
+// the batching rule.
+//
+// The batching rules are stored as kernels on the FuncTorchBatched key, so the way
+// VmapInterpreter calls the batching rule is roughly: (A) exclude all
+// dispatch keys aside from the Batched key, (B) redispatch so we get to the
+// Batched key.
+//
+// 2) sendToNextInterpreter(ophandle, stack)
+// The VmapInterpreter, when it sees aten::mv, will process it into a call to
+// aten::mm. It then needs to send the call to aten::mm to the next interpreter
+// in the interpreter stack.
+//
+// The VmapInterpreter just does this via a call to ophandle.callBoxed(stack)
+// and most Interpreters will implement it this way.
+
+enum class RandomnessType {
+    Error,      // always errors when calling a random function
+    Same,       // randomness appears the same across batches
+    Different,  // randomness appears different across batches
+    END
+};
+
+enum class TransformType {
+  Torch,  // Unused
+  Vmap,
+  Grad,  // reverse-mode AD, aka vjp
+  Jvp,  // forward-mode AD
+  Functionalize,
+};
+
+std::ostream& operator<<(std::ostream& os, const TransformType& t);
+
+// NOTE: [Interpreter "subclassing" design]
+//
+// How are various Interpreters for different transforms (vmap, grad, ...)
+// implemented?
+//
+// Accessing interpreters is in the hot-path of functorch so we have a constraint
+// that this code must be as fast as possible.
+//
+// As a result, we stay away from virtual methods and this causes our code
+// to look a little funny.
+//
+// `Interpreter` is the struct for Interpreters. It holds ALL of the
+// relevant information (what type of interpreter it is and the metadata).
+// Metadata for each interpreter is represented as a Union (std::variant)
+// of all possible metadata (VmapInterpreterMeta, GradInterpreterMeta, ...).
+//
+// Given an Interpreter, how do I get a "VmapInterpreter"? You may wish to do this
+// if you want to access the metadata fields (like batchSize and randomness).
+//
+// Each type of interpreter (e.g. Vmap) has a convenience struct
+// (e.g. VmapInterpreterPtr) associated with it.
+//
+// Construct the convenience struct with VmapInterpreterPtr(Interpreter*),
+// and then one can access methods on VmapInterpreterPtr like so:
+// >>> VmapInterpreterPtr(&interpreter).batchSize()
+//
+// Finally, Interpreter::process switches on the type of the interpreter
+// and calls one of {Transform}Intepreter::processImpl under the hood.
+// Same for Interpreter::sendToNextInterpreter :)
+
+struct VmapInterpreterMeta {
+  explicit VmapInterpreterMeta(c10::SymInt batchSize, RandomnessType randomness) :
+    batchSize_(std::move(batchSize)), randomness_(randomness) {}
+  c10::SymInt batchSize_;
+  RandomnessType randomness_;
+};
+
+struct GradInterpreterMeta {
+  explicit GradInterpreterMeta(bool prevGradMode): prevGradMode_(prevGradMode) {}
+  bool prevGradMode_;
+};
+
+struct JvpInterpreterMeta {
+  explicit JvpInterpreterMeta(bool prevFwdGradMode) : prevFwdGradMode_(prevFwdGradMode) {}
+  bool prevFwdGradMode_;
+};
+
+struct FunctionalizeInterpreterMeta {
+  explicit FunctionalizeInterpreterMeta(bool functionalizeAddBackViews) :
+    functionalizeAddBackViews_(functionalizeAddBackViews) {}
+  bool functionalizeAddBackViews_;
+};
+
+typedef std::variant<
+  int64_t,
+  GradInterpreterMeta,
+  JvpInterpreterMeta,
+  VmapInterpreterMeta,
+  FunctionalizeInterpreterMeta
+> InterpreterMeta;
+
+
+struct Interpreter {
+  // factory functions
+  static Interpreter Vmap(int64_t level, c10::SymInt batchSize, RandomnessType randomness) {
+    return Interpreter(TransformType::Vmap, level, VmapInterpreterMeta(std::move(batchSize), randomness));
+  }
+  static Interpreter Grad(int64_t level, bool prevGradMode) {
+    return Interpreter(TransformType::Grad, level, GradInterpreterMeta(prevGradMode));
+  }
+  static Interpreter Jvp(int64_t level, bool prevFwdGradMode) {
+    return Interpreter(TransformType::Jvp, level, JvpInterpreterMeta(prevFwdGradMode));
+  }
+  static Interpreter Functionalize(int64_t level, bool functionalizeAddBackViews) {
+    return Interpreter(TransformType::Functionalize, level, FunctionalizeInterpreterMeta(functionalizeAddBackViews));
+  }
+
+  // methods
+  TransformType key() const { return type_; }
+  int64_t level() const { return level_; }
+  const InterpreterMeta& meta() const { return meta_; }
+
+  void process(const c10::OperatorHandle& op, torch::jit::Stack* stack);
+  void sendToNextInterpreter(const c10::OperatorHandle& op, torch::jit::Stack* stack, bool grad_special_case);
+
+  void saveLocalDispatchKeySet(c10::impl::LocalDispatchKeySet keyset) {
+    TORCH_INTERNAL_ASSERT(!savedLocalDispatchKeySet_.has_value());
+    savedLocalDispatchKeySet_ = keyset;
+  }
+  void clearSavedLocalDispatchKeySet() {
+    TORCH_INTERNAL_ASSERT(savedLocalDispatchKeySet_.has_value());
+    savedLocalDispatchKeySet_ = c10::nullopt;
+  }
+  c10::impl::LocalDispatchKeySet getSavedLocalDispatchKeySet() const {
+    TORCH_INTERNAL_ASSERT(savedLocalDispatchKeySet_.has_value());
+    return *savedLocalDispatchKeySet_;
+  }
+
+  // An Interpreter is alive if we are currently inside the ongoing transform
+  // for the interpreter. For example, vmap(f)(x); inside of f, the vmap's
+  // corresponding Interpreter is alive, even when it is not on the DynamicLayerStack.
+  bool is_alive() const {
+    return *is_alive_;
+  }
+  const std::shared_ptr<bool>& is_alive_ptr() const {
+    return is_alive_;
+  }
+  void set_is_alive(bool alive) {
+    *is_alive_ = alive;
+  }
+
+  // Please don't use this
+  explicit Interpreter() = default;
+
+ private:
+  explicit Interpreter(TransformType type, int64_t level, InterpreterMeta meta):
+    type_(type), level_(level), is_alive_(std::make_shared<bool>(false)), meta_(std::move(meta)) {}
+
+  // fields
+  TransformType type_{};
+  int64_t level_{};
+  optional<c10::impl::LocalDispatchKeySet> savedLocalDispatchKeySet_;
+  std::shared_ptr<bool> is_alive_;
+  InterpreterMeta meta_;
+};
+
+// Applies the following for-loop:
+// for i in range(begin, end):
+//   args[i] = func(args[i])
+void foreachTensorInplace(std::vector<IValue>& args, int64_t begin, int64_t end,
+    std::function<Tensor(const Tensor&)> func);
+
+// Applies the following for-loop:
+// for i in range(begin, end):
+//   if use_flag_relative[i] == 1: <-- treats use_flag_relative as a bitset
+//     args[i] = func(args[i], i - begin, true)
+//   args[i] = func(args[i], i - begin)
+void foreachTensorInplaceWithFlag(std::vector<IValue>& args, int64_t begin, int64_t end,
+    const std::bitset<64> use_flag_relative, const std::function<Tensor(const Tensor&, bool)>& func);
+
+std::vector<int64_t> findUnwrappedInputs(std::vector<IValue>& args, int64_t begin, int64_t end);
+
+DispatchKeySet keysToExcludeWhenEnteringDynamicLayer(TransformType key);
+
+void setup_dispatch_key_tls(TransformType key, DispatchKeySet include);
+
+void sanityCheckStack(const c10::OperatorHandle& op, torch::jit::Stack* stack);
+
+} // namespace at::functorch

pythonProject/.venv/Lib/site-packages/torch/include/ATen/functorch/LegacyVmapTransforms.h

@@ -0,0 +1,187 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <ATen/functorch/Macros.h>
+#include <ATen/functorch/BatchedTensorImpl.h>
+
+namespace at::functorch {
+
+// This files contains the legacy (now-deprecated) batching rule API.
+// Please try to use the new-style batching rule API (see writing_batch_rules.md)
+
+// This file contains abstractions used for transforming *logical* vmap arguments
+// into *physical* arguments. (Keep reading for definitions of these terms).
+
+// NOTE: [Logical vs physical args]
+// Consider the following vmap.
+//   vmap(vmap(func, in_dims=(2,)), in_dims=(0,))(torch.ones(2, 3, 4))
+// This would produce a BatchedTensor wrapping a Tensor of size [2, 3, 4],
+// with batch dims 0 and 2:
+//   BatchedTensor(ones(2, 3, 4), bdims=[(lvl=1,dim=0),(lvl=2,dim=2)])
+//
+// We say the *logical* view of the tensor has size [3] -- tensors inside
+// `func` appear to have size [3].
+// However, the *physical* underlying tensor (the one passed to vmap) has size
+// [2, 3, 4].
+//
+// This notion of logical vs physical also extends to non-tensor arguments.
+// Consider the previous tensor; let's assume the user called
+// `torch.sum(tensor, dim=0)` inside of `func`. Then the logical
+// dimension they are reducing over is dim 0 but the physical dim is dim 1
+// (the first non-batch dimension)
+
+// Forward declared; see NOTE: [What is a VmapPhysicalView?]
+struct VmapPhysicalView;
+
+// Most PyTorch operators take 4 or fewer inputs.
+constexpr int64_t kVmapTransformStaticInputSize = 4;
+using VmapPhysicalViewVec = SmallVector<VmapPhysicalView, kVmapTransformStaticInputSize>;
+
+// Pytorch generally advertises good performance for <= 5 dims.
+// (see ATen/core/DimVector.h). We add a few extra dims (~3) for vmap
+// dimensions to get 8. Adjust this number as necessary
+constexpr int64_t kVmapStaticDimVecSize = 8;
+using VmapDimVector = SmallVector<int64_t, kVmapStaticDimVecSize>;
+using VmapSymDimVector = SmallVector<c10::SymInt, kVmapStaticDimVecSize>;
+
+// NOTE: [What is an VmapTransform?]
+// An *VmapTransform* converts logical views of tensors to physical views.
+//
+// Batching rules use VmapTransforms to convert logical arguments to
+// physical arguments, then call one or more at:: operator that handles the
+// physical arguments, and then converts the physical result back to a logical
+// argument.
+
+// VmapTransform for operators that take tensors with multiple batch dims.
+// Given one or more logical views on Tensors, `logicalToPhysical`
+// permutes all of the batch dims to the front of the tensor, aligns
+// and expands the batch dims to match each other (according to their `level`),
+// and returns a VmapPhysicalView on the tensor(s).
+struct TORCH_API MultiBatchVmapTransform {
+  static VmapPhysicalView logicalToPhysical(const Tensor& logical_tensor);
+  static VmapPhysicalViewVec logicalToPhysical(ITensorListRef logical_tensors);
+};
+
+// VmapTransform for operators that broadcast all inputs.
+// Given some logical views on Tensors, `logicalToPhysical`:
+// - permutes all of the batch dims to the front of the tensors
+// - aligns all the batch dims to the collective levels of all of the tensors.
+//   If a tensor does not have a batch dim for a vmap level, then it receives
+//   a size-one dimension for said level.
+// - aligns the non-batch dims to have the same dimensionality, adding extra
+//   size-1 dimensions in between the batch dimensions and the non-batch dimensions
+//   so that the batch dimensions are lined up from the right.
+//
+// For example: given inputs of size (B, 2) and (B, 3, 2) where B is the batch
+// dimension, BroadcastingVmapTransform returns VmapPhysicalViews that wrap tensors
+// of size (B, 1, 2) and (B, 3, 2).
+//
+// Given inputs of size (B, 2) and (2,), BroadcastingVmapTransform returns
+// VmapPhysicalViews wrapping tensors of size (B, 2) and (1, 2). We don't
+// actually *need* to return a tensor of size (1, 2) for the second tensor
+// because the broadcasting operation takes care of that for us, but we do
+// it anyways to keep things simple.
+struct TORCH_API BroadcastingVmapTransform {
+  static VmapPhysicalViewVec logicalToPhysical(TensorList logical_tensors);
+};
+
+// Forward declared, if you're reading this file head to toe, don't worry about
+// it yet.
+struct VmapPhysicalToLogicalMap;
+
+// NOTE: [What is a VmapPhysicalView?]
+// VmapPhysicalView represents a physical view on a Tensor.
+//
+// One can use it to further convert logical dimension indices, logical shapes,
+// and more to their physical variants, or convert a new (physical) tensor into
+// a logical BatchedTensor. (TODO(rzou): some of these are not yet implemented).
+//
+// VmapPhysicalView stores a physical tensor with all of its batch dimensions at
+// the front and some levels that correspond to said batch dimensions.
+//
+// The levels bitset specifies which vmap levels correspond to the batch
+// dimensions at the front of the tensor. In particular, the number of set bits
+// corresponds to the number of batch dimensions on `tensor` and the rightmost
+// bit of `levels` specifies the maximum number of nested vmaps we are in at
+// this point in time.
+// For example, given:
+//   physical_view = VmapPhysicalView(tensor=ones(2, 3, 4, 5, 6), levels={1, 3})
+//
+// Rightmost bit of `levels` is 3 indicating the number of nested vmaps less
+// than or equal to 3.
+//   bitset: 010100
+//              ^
+//              |
+//   levels: 012345
+struct TORCH_API VmapPhysicalView {
+  VmapPhysicalView(Tensor&& tensor, std::bitset<kVmapNumLevels> levels)
+      : levels_(levels), tensor_(std::move(tensor)) {
+    // TORCH_INTERNAL_ASSERT(!isBatchedTensor(tensor));
+  }
+
+  Tensor& tensor() { return tensor_; }
+  const Tensor& tensor() const { return tensor_; }
+
+  // Maps logical dim indices to physical dim indices. Also does dim wrapping.
+  //
+  // For example, given:
+  //   physical_view = VmapPhysicalView(tensor=ones(2, 3, 4, 5), levels={1, 3})
+  //
+  // Then physical_view.getPhysicalDims({0, 1}) returns {2, 3}.
+  // This is because the size of levels tell us that the first two dimensions
+  // of `tensor_` are batch dimensions, so a logical dim of `n` is actually
+  // a physical dim of `n + 2`.
+  VmapDimVector getPhysicalDims(IntArrayRef logical_dims) const;
+  int64_t getPhysicalDim(int64_t logical_dim) const;
+
+  // Returns a VmapPhysicalToLogicalMap object. This can be used for
+  // mapping a physical tensor to a new logical tensor (BatchedTensor)
+  VmapPhysicalToLogicalMap getPhysicalToLogicalMap() const;
+
+  // Maps a logical shape to a physical shape by pre-pending the batch
+  // sizes to the logical shape.
+  VmapDimVector getPhysicalShape(IntArrayRef logical_shape) const;
+  SymDimVector getPhysicalShape(c10::SymIntArrayRef logical_shape) const;
+
+  int64_t numBatchDims() const;
+
+ private:
+  int64_t numLogicalDims() const;
+
+  std::bitset<kVmapNumLevels> levels_;
+  Tensor tensor_;
+};
+
+// Convenience struct used for mapping a physical tensor (a non-BatchedTensor)
+// to a logical one (BatchedTensor). It holds some levels that are used to do the
+// mapping and assumes that the batch dimensions in the physical tensor all
+// occur at the front of the tensor.
+struct TORCH_API VmapPhysicalToLogicalMap {
+  VmapPhysicalToLogicalMap(std::bitset<kVmapNumLevels> levels): levels_(levels) {}
+
+  // Maps a physical tensor to a new logical tensor (BatchedTensor).
+  // Assumes that all of the "batch dimensions" are at the front
+  // of the physical tensor. For example, given:
+  // - x = rank-4 Tensor with size 2, 3, 5, 7
+  // - levels = (2, 4)
+  // Returns:
+  // - BatchedTensor(x, bdims=[(dim=0,lvl=2), (dim=1, lvl=4)])
+  Tensor apply(const Tensor& physical_tensor) const;
+
+  // Given a vector of physical tensors,
+  // 1. maps each tensor to a new logical tensor. Assumes that all of the
+  //    "batch dimensions" are at the front of the physical tensors.
+  // 2. stores the new logical tensors back into the passed-in vector. This is
+  //    to avoid additional dynamic allocations.
+  void applyInplace(std::vector<Tensor>& physical_tensors) const;
+
+  std::bitset<kVmapNumLevels> levels_;
+};
+
+
+} // namespace at::functorch

pythonProject/.venv/Lib/site-packages/torch/include/ATen/functorch/Macros.h

@@ -0,0 +1,3 @@
+#pragma once
+
+#define SINGLE_ARG(...) __VA_ARGS__