James Mochizuki-Freeman

## Problem Description The LSTM weight matrices in `ppo_atari_lstm.py` seem to be be initialized incorrectly, if the goal is to have a separate orthogonal matrix for each gate. Since `lstm.weight_ih_l0`...

Downstream packages will expect the `py.typed` marker to be present, which explicitly declares that a library supports typing (see [relevant portion of typing spec](https://typing.python.org/en/latest/spec/distributing.html#packaging-type-information)). Pyright currently emits a [`reportMissingTypeStubs`](https://github.com/microsoft/pylance-release/blob/main/docs/diagnostics/reportMissingTypeStubs.md) diagnostic...

The `named_axis` functions `full()`, `zeros()` and `ones()` annotate their dtype parameters with the non-existent `np.DTypeLike`. This change replaces the annotations with `jax.typing.DTypeLike | None` to match the dtype parameter of...

Docstrings for reexported functions do not fall back to implementation definition

2

### Summary When hovering a reexported function from a module with a `.pyi` stub file, the docstring from the stub file definition is (correctly) shown instead of the implementation docstring....

Built-in operators on `JustFloat`s should return `JustFloat`.

1

### Summary (Based on https://github.com/astral-sh/ty/issues/2184#issuecomment-3697621323) Ty reports the sum of two `JustFloat` values is `int | float`. I would have expected it to be another `JustFloat`, given that result is...