pytensor 3.2.4__tar.gz → 3.3.0__tar.gz

{pytensor-3.2.4/pytensor.egg-info → pytensor-3.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.2.4
+Version: 3.3.0
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause
@@ -24,6 +24,7 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
+Classifier: Programming Language :: Python :: Free Threading :: 1 - Unstable
 Requires-Python: <3.15,>=3.12
 Description-Content-Type: text/x-rst
 License-File: LICENSE.txt

{pytensor-3.2.4 → pytensor-3.3.0}/pyproject.toml

@@ -33,6 +33,7 @@ classifiers = [
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
     "Programming Language :: Python :: 3.14",
+    "Programming Language :: Python :: Free Threading :: 1 - Unstable",
 ]
 
 keywords = [
@@ -185,11 +186,18 @@ files = ["pytensor", "tests"]
 build = "*"
 # Uncomment to skip builds that compile but fail when trying to test (maybe due to incompatibility with runner)
 # archs = ["auto64"]
-# Disable any-platform (pp*), 32-bit builds, and free-threaded builds (cp*t-*)
+# Disable any-platform (pp*) and 32-bit builds
 # Additional options to consider: "*musllinux*"
-skip = ["pp*", "*-win32", "*-manylinux_i686", "cp*t-*"]
+skip = ["pp*", "*-win32", "*-manylinux_i686"]
+# Free-threaded (cp*t-*) wheels are built by default. The compiled extensions
+# are not marked free-threading-safe, so importing them re-enables the GIL
+# (with a warning). That's fine: they're no longer loaded by default.
 build-frontend = "build"
-test-command = 'python -c "import pytensor; print(pytensor.__version__); from pytensor.scan import scan_perform; print(scan_perform.get_version())"'
+# Import the compiled extensions and evaluate a function to check the wheel loads and
+# runs everywhere, including free-threaded builds. cxx="" disables the C backend (the
+# minimal wheel-test env has no libpython to link a runtime-compiled thunk against) and
+# FAST_COMPILE avoids numba, which crashes on free-threaded Windows.
+test-command = '''python -c "import pytensor; pytensor.config.cxx = ''; import pytensor.tensor as pt; print(pytensor.__version__); from pytensor.scan import scan_perform; print(scan_perform.get_version()); x = pt.vector(); f = pytensor.function([x], (x * 2).sum(), mode='FAST_COMPILE'); assert float(f([1.0, 2.0, 3.0])) == 12.0"'''
 test-skip = ["*musllinux*", "*i686*"]
 
 # Testing seems to be running into issues locating libs where expected

{pytensor-3.2.4 → pytensor-3.3.0}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-08-02T00:12:33+0200",
+ "date": "2026-08-12T10:16:27+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "7175f7cadf3934a03f63e6e7b187fa1519e6d85f",
- "version": "3.2.4"
+ "full-revisionid": "aa1bc7772eb11ed8def2814374aeb5da03da590c",
+ "version": "3.3.0"
 }
 '''  # END VERSION_JSON
 

{pytensor-3.2.4 → pytensor-3.3.0}/pytensor/link/jax/dispatch/subtensor.py

@@ -1,3 +1,5 @@
+import jax.numpy as jnp
+
 from pytensor.link.jax.dispatch.basic import jax_funcify
 from pytensor.tensor.subtensor import (
     AdvancedIncSubtensor,
@@ -60,7 +62,11 @@ def jax_funcify_IncSubtensor(op, node, **kwargs):
             indices = indices[0]
 
         if isinstance(op, AdvancedIncSubtensor):
-            op._check_runtime_broadcast_of_vector_index(node, x, y, indices)
+            # jax_typify downgrades 0d arrays to Python scalars, which have no .shape,
+            # so re-arrayify y for a check that is written against array values.
+            op._check_runtime_broadcast_of_vector_index(
+                node, x, jnp.asarray(y), indices
+            )
 
         return jax_fn(x, indices, y)
 

{pytensor-3.2.4 → pytensor-3.3.0}/pytensor/link/mlx/dispatch/basic.py

@@ -15,25 +15,36 @@ from pytensor.link.utils import fgraph_to_python
 from pytensor.raise_op import Assert, CheckAndRaise
 
 
+def float64_supported():
+    """Return whether the current default device can operate on float64.
+
+    MLX implements float64 on the CPU only; the Metal backend rejects it outright.
+    """
+    return mx.default_device() == mx.cpu
+
+
 def convert_dtype_to_mlx(dtype_str, auto_cast_unsupported=True):
     """Convert PyTensor dtype strings to MLX dtype objects.
 
     MLX expects dtype objects rather than string literals for type conversion.
     This function maps common dtype strings to their MLX equivalents.
 
+    float64 survives when the default device supports it and is narrowed to float32
+    otherwise. complex128 is always narrowed, MLX having no wider complex type.
+
     Parameters
     ----------
     dtype_str : str or MLX dtype
         The dtype to convert
-    auto_cast_unsupported : bool
-        If True, automatically cast unsupported dtypes to supported ones with warnings
+    auto_cast_unsupported : bool, optional
+        If True, narrow dtypes the current device cannot handle, warning as it does so.
+        If False, return the requested dtype and let any later failure surface. Default
+        True.
 
     Returns
     -------
     MLX dtype object
     """
-    import warnings
-
     if isinstance(dtype_str, str):
         if dtype_str == "bool":
             return mx.bool_
@@ -58,11 +69,12 @@ def convert_dtype_to_mlx(dtype_str, auto_cast_unsupported=True):
         elif dtype_str == "float32":
             return mx.float32
         elif dtype_str == "float64":
-            if auto_cast_unsupported:
+            if auto_cast_unsupported and not float64_supported():
                 warnings.warn(
                     "MLX does not support float64 on GPU. Automatically casting to float32. "
-                    "This may result in reduced precision. To avoid this warning, "
-                    "explicitly use float32 in your code or set floatX='float32' in PyTensor config.",
+                    "This may result in reduced precision. To keep float64, run on the CPU "
+                    "device with mx.set_default_device(mx.cpu); to avoid this warning, use "
+                    "float32 or set floatX='float32' in PyTensor config.",
                     UserWarning,
                     stacklevel=3,
                 )
@@ -116,6 +128,10 @@ def mlx_typify(data, **kwargs):
 
 @mlx_typify.register(np.ndarray)
 def mlx_typify_tensor(data, dtype=None, **kwargs):
+    # mx.array narrows float64 input to float32 unless the dtype is named explicitly,
+    # and it does so on the CPU too, where float64 is perfectly usable
+    if dtype is None and data.dtype == np.float64 and float64_supported():
+        dtype = mx.float64
     return _nan_safe_constant(data, dtype=dtype)
 
 

pytensor-3.3.0/pytensor/link/mlx/dispatch/scalar/__init__.py

@@ -0,0 +1 @@
+from pytensor.link.mlx.dispatch.scalar import basic, bessel, erf, erfcinv, gamma, math

pytensor-3.2.4/pytensor/link/mlx/dispatch/scalar.py → pytensor-3.3.0/pytensor/link/mlx/dispatch/scalar/basic.py

@@ -5,16 +5,15 @@ import mlx.core as mx
 from pytensor.link.mlx.dispatch.basic import convert_dtype_to_mlx, mlx_funcify
 from pytensor.scalar.basic import (
     Cast,
+    Clip,
     Composite,
     Identity,
     Mod,
     ScalarOp,
     Second,
 )
-from pytensor.scalar.math import Erfc, Erfcx, Log1mexp, Sigmoid, Softplus
 
 
-# MLX name overrides for nfunc_spec names that don't match mlx.core
 MLX_NFUNC_OVERRIDES = {
     "true_divide": "divide",
     "invert": "bitwise_invert",
@@ -96,6 +95,14 @@ def mlx_funcify_Mod(op, **kwargs):
     return mlx_mod
 
 
+@mlx_funcify.register(Clip)
+def mlx_funcify_Clip(op, **kwargs):
+    def clip(x, min, max):
+        return mx.where(x < min, min, mx.where(x > max, max, x))
+
+    return clip
+
+
 @mlx_funcify.register(Identity)
 def mlx_funcify_Identity(op, **kwargs):
     def identity(x):
@@ -115,55 +122,6 @@ def mlx_funcify_Second(op, **kwargs):
     return second
 
 
-@mlx_funcify.register(Sigmoid)
-def mlx_funcify_Sigmoid(op, **kwargs):
-    return mx.sigmoid
-
-
-@mlx_funcify.register(Erfc)
-def mlx_funcify_Erfc(op, **kwargs):
-    def erfc(x):
-        return 1.0 - mx.erf(x)
-
-    return erfc
-
-
-@mlx_funcify.register(Erfcx)
-def mlx_funcify_Erfcx(op, **kwargs):
-    def erfcx(x):
-        return mx.exp(x * x) * (1.0 - mx.erf(x))
-
-    return erfcx
-
-
-@mlx_funcify.register(Softplus)
-def mlx_funcify_Softplus(op, **kwargs):
-    def softplus(x):
-        return mx.where(
-            x < -37.0,
-            mx.exp(x),
-            mx.where(
-                x < 18.0,
-                mx.log1p(mx.exp(x)),
-                mx.where(
-                    x < 33.3,
-                    x + mx.exp(-x),
-                    x,
-                ),
-            ),
-        )
-
-    return softplus
-
-
-@mlx_funcify.register(Log1mexp)
-def mlx_funcify_Log1mexp(op, node, **kwargs):
-    def log1mexp(x):
-        return mx.where(x < mx.log(0.5), mx.log1p(-mx.exp(x)), mx.log(-mx.expm1(x)))
-
-    return log1mexp
-
-
 @mlx_funcify.register(Composite)
 def mlx_funcify_Composite(op, node=None, **kwargs):
     return mlx_funcify(op.fgraph, squeeze_output=True)