🔒 Enforcing Type Inside Class Method

[lilac spruce]

I have a Tensor class:

from __future__ import annotations
from typing import Any, List, TYPE_CHECKING
import numbers

import jax
import jax.numpy as jnp
from jax.typing import DTypeLike, ArrayLike
import numpy as np

# bunch of other imports

class Tensor:
    def __init__(
        self,
        data: ArrayLike,
        name: str | None = None,
        dtype: DTypeLike | None = None,
        device: Any | None = None,
        requires_grad: bool = False,
    ) -> None:
        self.data: jax.Array = jax.device_put(
            jnp.array(data, dtype=dtype), device=device
        )

        if requires_grad and not jnp.issubdtype(self.data.dtype, jnp.floating):
            raise ValueError(
                f"Only floating-point dtypes can have requires_grad=True. "
                f"Got dtype={self.data.dtype}. Convert to float first."
            )

        self.name: str = name or ""
        self.grads: jax.Array | None = None
        self.grad_fn: GradFunction | None = None
        self.requires_grad: bool = requires_grad

[rough monolith]

@lilac spruce

Python help channel opened

Remember to:

• Ask your Python question, not if you can ask or if there's an expert who can help.
• Show a code sample as text (rather than a screenshot) and the error message, if you've got one.
• Explain what you expect to happen and what actually happens.

:warning: Do not pip install anything that isn't related to your question, especially if asked to over DMs.

Closes after a period of inactivity, or when you send !close.

[lilac spruce]

def __add__(self, other_tensor: TensorLike) -> Tensor:
        """
        Element-wise addition with automatic broadcasting and gradient tracking.

        Parameters
        ----------
        other_tensor : TensorLike
            Value to add. Accepts tensors, scalars, bools, JAX arrays, or NumPy arrays.
            Bools are treated as 0 (False) or 1 (True). Complex numbers are not supported.

        Returns
        -------
        Tensor
            Result with shape following NumPy broadcasting rules: https://numpy.org/doc/stable/user/basics.broadcasting.html

        Raises
        ------
        TypeError
            If other_tensor is not a supported type.

        Notes
        -----
        Gradient computation: For :math:`z = a + b`, :math:`dz/da = 1`, :math:`dz/db = 1`.
        """
        if not isinstance(other_tensor, (int, float, bool, jax.Array, Tensor)):
            raise TypeError(
                f"Unsupported type for addition with Tensor. Can't add {type(other_tensor)} to Tensor."
            )

        if isinstance(other_tensor, jax.Array):
            other_tensor = Tensor(other_tensor)

        if isinstance(other_tensor, (int, float, bool)):
            other_tensor = Tensor(jnp.array(other_tensor))

        data: jax.Array = self.data + other_tensor.data

        ctx = Context()
        ctx.save_for_backward(self.data, other_tensor.data)

        result = Tensor(
            data=data,
            dtype=data.dtype,
            device=data.device,
            requires_grad=self.requires_grad or other_tensor.requires_grad,
        )

        result.grad_fn = GradFunction(
            backward_fn=add_backward,
            ctx=ctx,
            parents=[self, other_tensor],
        )

        return result

    # bunch of other defs

TensorLike = Tensor | jax.Array | np.ndarray | bool | int | float | numbers.Number
tensor = Tensor

which basically isa glorified JAX wrapper but implementing its own auto diff system to enable deep learning research in that field.

Anyway, You can see that inside my methods like e.g. Tensor::__add__ I do a bunch of type checking to make sure I work with a Tensor. You can argue it's useless since I do mathematical operations on JAX objects but that'd just invert the problem i.e. instead of having a bunch of type checking for creating a Node, I'd have it for getting .data.

the current approach also has the problem that certain types are subclasses of each other (which I only just learned). e.g. issubclass(bool, int) is True.

Any suggestions on how to implement this nicer?

[rough monolith]

@lilac spruce

Python help channel closed for inactivity

This help channel has been closed. Feel free to create a new post in #1035199133436354600. To maximize your chances of getting a response, check out this guide on asking good questions.