Skip to content

Default typify in JITLinkers expects tensors/arrays, why? #819

Open
@ricardoV94

Description

@ricardoV94

Description

@singledispatch
def jax_typify(data, dtype=None, **kwargs):
r"""Convert instances of PyTensor `Type`\s to JAX types."""
if dtype is None:
return data
else:
return jnp.array(data, dtype=dtype)

(And now from copying the approach also the PyTorch Linker)

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions