mirror of
https://github.com/wassname/jaxtyping.git
synced 2026-08-21 11:16:29 +08:00
Roman Knyazhitskiy
172b83b4fc
* Add a test for generators * Remove output annotations from decorators Also guarded torch imports for better compatibility with requirements.txt * Add flag to the main meta class to skip the typecheck * Return to the old solution * Make async tests work * Minor adjustments/fixing typos * Correct Python path for new tests * Remove some jax-dependent code * Implement equality for MetaArrays * Make all Dim variations frozen dataclasses * Shorten AbstractArray methods * Final touches * Removing get_origin use * Update tests with @jaxtyp
36 lines
838 B
Python
36 lines
838 B
Python
from typing import Tuple, Union
|
|
|
|
import pytest
|
|
|
|
from jaxtyping import (
|
|
Array,
|
|
Float,
|
|
Float32,
|
|
Integer,
|
|
PRNGKeyArray,
|
|
PyTree,
|
|
Shaped,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"make_fn",
|
|
[
|
|
lambda: Float[Array, "4"],
|
|
lambda: Float32[Array, ""],
|
|
lambda: Integer[Array, "1 2 3"],
|
|
lambda: Shaped[PRNGKeyArray, "2"],
|
|
lambda: Float[float, "#*shape"],
|
|
lambda: PyTree[int],
|
|
lambda: PyTree[Float[Array, ""]],
|
|
lambda: PyTree[Float32[Array, "*m b c"]],
|
|
lambda: PyTree[PyTree[Float32[Array, "1 2 b *"]]],
|
|
lambda: PyTree[Union[str, Float32[Array, "1"]]],
|
|
lambda: PyTree[
|
|
Tuple[int, float, Float[Array, ""], PyTree[Union[Float[Array, ""], float]]]
|
|
],
|
|
],
|
|
)
|
|
def test_equals(make_fn):
|
|
assert make_fn() == make_fn()
|