Skip to content

MLX backend: every pad raises TypeError (pad_width passed as mx.array) #2392

Description

@guillaume-osmo

Description

mlx_funcify_pad forwards pad_width straight to mx.pad, but the linker typifies every input to mx.array while mx.pad takes an int or a list of (before, after) int pairs, so no pad works on this backend. Same root cause as #2386.

test_mlx_pad[constant_default] and test_mlx_pad[edge] are already red on main.

import numpy as np
import pytensor
import pytensor.tensor as pt

x = pt.matrix("x", shape=(3, 4), dtype="float32")
xv = np.zeros((3, 4), dtype="float32")

print(pytensor.function([x], pt.pad(x, pad_width=2, mode="edge"), mode="CVM")(xv).shape)  # (7, 8)
print(pytensor.function([x], pt.pad(x, pad_width=2, mode="edge"), mode="MLX")(xv).shape)
# TypeError: pad(): incompatible function arguments
#   pad_width: int | tuple[int] | tuple[int, int] | list[tuple[int, int]]

Note the value cannot simply be read back at runtime: the linker enables mx.compile by default and MLX forbids evaluating a traced array, so a constant pad_width has to be resolved at funcify time, as in #2386.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions