Fix merge conflicts

Author: HarshvirSandhu

.github/workflows/pypi.yml

@@ -30,12 +30,13 @@ jobs:
       - name: Build SDist
         run: pipx run build --sdist
 
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
         with:
+          name: sdist
           path: dist/*.tar.gz
 
   build_wheels:
-    name: Build ${{ matrix.python-version }} wheels on ${{ matrix.platform }}
+    name: Build wheels for ${{ matrix.platform }}
     runs-on: ${{ matrix.platform }}
     strategy:
       matrix:
@@ -51,19 +52,27 @@ jobs:
       - name: Build wheels
         uses: pypa/cibuildwheel@v2.19.2
 
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
         with:
+          name: wheels-${{ matrix.platform }}
           path: ./wheelhouse/*.whl
 
   check_dist:
     name: Check dist
     needs: [make_sdist,build_wheels]
     runs-on: ubuntu-22.04
     steps:
-      - uses: actions/download-artifact@v3
+      - uses: actions/download-artifact@v4
         with:
-          name: artifact
+          name: sdist
           path: dist
+
+      - uses: actions/download-artifact@v4
+        with:
+          pattern: wheels-*
+          path: dist
+          merge-multiple: true
+
       - name: Check SDist
         run: |
           mkdir -p test-sdist
@@ -83,12 +92,18 @@ jobs:
     runs-on: ubuntu-latest
     if: github.event_name == 'release' && github.event.action == 'published'
     steps:
-    - uses: actions/download-artifact@v3
-      with:
-        name: artifact
-        path: dist
+      - uses: actions/download-artifact@v4
+        with:
+          name: sdist
+          path: dist
 
-    - uses: pypa/gh-action-pypi-publish@v1.9.0
-      with:
-        user: __token__
-        password: ${{ secrets.pypi_password }}
+      - uses: actions/download-artifact@v4
+        with:
+          pattern: wheels-*
+          path: dist
+          merge-multiple: true
+
+      - uses: pypa/gh-action-pypi-publish@v1.9.0
+        with:
+          user: __token__
+          password: ${{ secrets.pypi_password }}

.github/workflows/test.yml

@@ -187,9 +187,9 @@ jobs:
           FLOAT32: ${{ matrix.float32 }}
 
       - name: Upload coverage file
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
-          name: coverage
+          name: coverage-${{ steps.matrix-id.outputs.id }}
           path: coverage/coverage-${{ steps.matrix-id.outputs.id }}.xml
 
   benchmarks:
@@ -273,10 +273,11 @@ jobs:
           python -m pip install -U coverage>=5.1 coveralls
 
       - name: Download coverage file
-        uses: actions/download-artifact@v3
+        uses: actions/download-artifact@v4
         with:
-          name: coverage
+          pattern: coverage-*
           path: coverage
+          merge-multiple: true
 
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v4

pytensor/compile/builders.py

@@ -8,7 +8,6 @@
 
 from pytensor.compile.function import function
 from pytensor.compile.function.pfunc import rebuild_collect_shared
-from pytensor.compile.mode import optdb
 from pytensor.compile.sharedvalue import SharedVariable
 from pytensor.configdefaults import config
 from pytensor.gradient import DisconnectedType, Rop, grad
@@ -24,7 +23,6 @@
 from pytensor.graph.null_type import NullType
 from pytensor.graph.op import HasInnerGraph, Op
 from pytensor.graph.replace import clone_replace
-from pytensor.graph.rewriting.basic import in2out, node_rewriter
 from pytensor.graph.utils import MissingInputError
 
 
@@ -575,7 +573,7 @@ def lop_overrides(inps, grads):
             for inp_grad in input_grads
             if not isinstance(inp_grad.type, DisconnectedType | NullType)
         ]
-        lop_op = type(self)(
+        lop_op = OpFromGraph(
             inputs=inner_inputs + connected_inner_outputs + connected_output_grads,
             outputs=connected_input_grads,
             inline=self.is_inline,
@@ -669,7 +667,7 @@ def _build_and_cache_rop_op(self):
             for out_grad in output_grads
             if not isinstance(out_grad.type, DisconnectedType | NullType)
         ]
-        rop_op = type(self)(
+        rop_op = OpFromGraph(
             inputs=inner_inputs + eval_points,
             outputs=filtered_output_grads,
             inline=self.is_inline,
@@ -852,29 +850,3 @@ def perform(self, node, inputs, outputs):
         assert len(variables) == len(outputs)
         for output, variable in zip(outputs, variables):
             output[0] = variable
-
-
-@node_rewriter([OpFromGraph])
-def inline_ofg_expansion(fgraph, node):
-    """
-    This optimization expands internal graph of OpFromGraph.
-    Only performed if node.op.is_inline == True
-    Doing so can improve optimization at the cost of compilation speed.
-    """
-    op = node.op
-    if not isinstance(op, OpFromGraph):
-        return False
-    if not op.is_inline:
-        return False
-    return clone_replace(op.inner_outputs, dict(zip(op.inner_inputs, node.inputs)))
-
-
-# We want to run this before the first merge optimizer
-# and before the first scan optimizer.
-optdb.register(
-    "inline_ofg_expansion",
-    in2out(inline_ofg_expansion),
-    "fast_compile",
-    "fast_run",
-    position=-0.01,
-)

pytensor/link/__init__.py

@@ -0,0 +1 @@
+from pytensor.link.pytorch.linker import PytorchLinker

pytensor/link/jax/dispatch/__init__.py

@@ -6,6 +6,7 @@
 import pytensor.link.jax.dispatch.blockwise
 import pytensor.link.jax.dispatch.elemwise
 import pytensor.link.jax.dispatch.extra_ops
+import pytensor.link.jax.dispatch.pad
 import pytensor.link.jax.dispatch.math
 import pytensor.link.jax.dispatch.nlinalg
 import pytensor.link.jax.dispatch.random

pytensor/link/jax/dispatch/basic.py

@@ -1,10 +1,13 @@
 import warnings
+from collections.abc import Callable
 from functools import singledispatch
 
 import jax
 import jax.numpy as jnp
 import numpy as np
 
+from pytensor.compile import JAX
+from pytensor.compile.builders import OpFromGraph
 from pytensor.compile.ops import DeepCopyOp, ViewOp
 from pytensor.configdefaults import config
 from pytensor.graph.fg import FunctionGraph
@@ -114,3 +117,24 @@ def viewop(x):
         return x
 
     return viewop
+
+
+@jax_funcify.register(OpFromGraph)
+def jax_funcify_OpFromGraph(ofg: OpFromGraph, node=None, **kwargs) -> Callable:
+    _ = kwargs.pop("storage_map", None)
+
+    # Apply inner rewrites
+    JAX.optimizer(ofg.fgraph)
+    fgraph_fn = jax_funcify(ofg.fgraph, **kwargs)
+
+    if len(ofg.fgraph.outputs) == 1:
+
+        def opfromgraph(*inputs):
+            return fgraph_fn(*inputs)[0]
+
+    else:
+
+        def opfromgraph(*inputs):
+            return fgraph_fn(*inputs)
+
+    return opfromgraph

pytensor/link/jax/dispatch/pad.py

@@ -0,0 +1,53 @@
+import jax.numpy as jnp
+import numpy as np
+
+from pytensor.link.jax.dispatch import jax_funcify
+from pytensor.tensor.pad import Pad
+
+
+@jax_funcify.register(Pad)
+def jax_funcify_pad(op, **kwargs):
+    pad_mode = op.pad_mode
+    reflect_type = op.reflect_type
+    has_stat_length = op.has_stat_length
+
+    if pad_mode == "constant":
+
+        def constant_pad(x, pad_width, constant_values):
+            return jnp.pad(x, pad_width, mode=pad_mode, constant_values=constant_values)
+
+        return constant_pad
+
+    elif pad_mode == "linear_ramp":
+
+        def lr_pad(x, pad_width, end_values):
+            # JAX does not allow a dynamic input if end_values is non-scalar
+            if not isinstance(end_values, int | float):
+                end_values = tuple(np.array(end_values))
+            return jnp.pad(x, pad_width, mode=pad_mode, end_values=end_values)
+
+        return lr_pad
+
+    elif pad_mode in ["maximum", "minimum", "mean"] and has_stat_length:
+
+        def stat_pad(x, pad_width, stat_length):
+            # JAX does not allow a dynamic input here, need to cast to tuple
+            return jnp.pad(
+                x, pad_width, mode=pad_mode, stat_length=tuple(np.array(stat_length))
+            )
+
+        return stat_pad
+
+    elif pad_mode in ["reflect", "symmetric"]:
+
+        def loop_pad(x, pad_width):
+            return jnp.pad(x, pad_width, mode=pad_mode, reflect_type=reflect_type)
+
+        return loop_pad
+
+    else:
+
+        def pad(x, pad_width):
+            return jnp.pad(x, pad_width, mode=pad_mode)
+
+        return pad

pytensor/link/pytorch/dispatch/__init__.py

@@ -2,8 +2,10 @@
 from pytensor.link.pytorch.dispatch.basic import pytorch_funcify, pytorch_typify
 
 # # Load dispatch specializations
+import pytensor.link.pytorch.dispatch.blas
 import pytensor.link.pytorch.dispatch.scalar
 import pytensor.link.pytorch.dispatch.elemwise
+import pytensor.link.pytorch.dispatch.math
 import pytensor.link.pytorch.dispatch.extra_ops
 import pytensor.link.pytorch.dispatch.subtensor
 import pytensor.link.pytorch.dispatch.sort

pytensor/link/pytorch/dispatch/blas.py

@@ -0,0 +1,14 @@
+import torch
+
+from pytensor.link.pytorch.dispatch import pytorch_funcify
+from pytensor.tensor.blas import BatchedDot
+
+
+@pytorch_funcify.register(BatchedDot)
+def pytorch_funcify_BatchedDot(op, **kwargs):
+    def batched_dot(a, b):
+        if a.shape[0] != b.shape[0]:
+            raise TypeError("Shapes must match in the 0-th dimension")
+        return torch.bmm(a, b)
+
+    return batched_dot

pytensor/link/pytorch/dispatch/math.py

@@ -0,0 +1,12 @@
+import torch
+
+from pytensor.link.pytorch.dispatch import pytorch_funcify
+from pytensor.tensor.math import Dot
+
+
+@pytorch_funcify.register(Dot)
+def pytorch_funcify_Dot(op, **kwargs):
+    def dot(x, y):
+        return torch.matmul(x, y)
+
+    return dot

pytensor/tensor/__init__.py

@@ -130,6 +130,7 @@ def _get_vector_length_Constant(op: Op | Variable, var: Constant) -> int:
 from pytensor.tensor.extra_ops import *
 from pytensor.tensor.io import *
 from pytensor.tensor.math import *
+from pytensor.tensor.pad import pad
 from pytensor.tensor.shape import (
     reshape,
     shape,