Skip to content

[tensor-shapes] add missing primitive declarations in jax.lax #5620

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
jakevdp wants to merge 1 commit into from
+165 −0
Closed

[tensor-shapes] add missing primitive declarations in jax.lax #5620

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
162 changes: 162 additions & 0 deletions tensor-shapes/jax-pyrefly-stubs/jax-stubs/lax/__init__.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@ from jax._src.sharding_impls import (
NamedSharding as _NamedSharding,
PartitionSpec as _PartitionSpec,
)
from jax.extend.core import Primitive as _Primitive
from jax.typing import DTypeLike
from shape_extensions import (
Flag,
Expand Down Expand Up @@ -1816,3 +1817,164 @@ def shape_as_value(shape: Sequence[int]) -> _Array[IntTuple]: ...
def stage[Shape: _Shape = []](x: _ArrayLike[Shape], /) -> _Array[Shape]: ...
def stop_gradient[T](x: T) -> T: ...
def with_sharding_constraint[T](x: T, shardings: Any) -> T: ...

# -----------------------------------------------------------------------------
# Primitives
# -----------------------------------------------------------------------------

abs_p: _Primitive
acos_p: _Primitive
acosh_p: _Primitive
add_p: _Primitive
after_all_p: _Primitive
all_gather_p: _Primitive
all_to_all_p: _Primitive
and_p: _Primitive
approx_top_k_p: _Primitive
argmax_p: _Primitive
argmin_p: _Primitive
asin_p: _Primitive
asinh_p: _Primitive
atan2_p: _Primitive
atan_p: _Primitive
atanh_p: _Primitive
axis_index_p: _Primitive
bessel_i0e_p: _Primitive
bessel_i1e_p: _Primitive
bitcast_convert_type_p: _Primitive
broadcast_in_dim_p: _Primitive
cbrt_p: _Primitive
ceil_p: _Primitive
clamp_p: _Primitive
clz_p: _Primitive
complex_p: _Primitive
concatenate_p: _Primitive
cond_p: _Primitive
conj_p: _Primitive
convert_element_type_p: _Primitive
conv_general_dilated_p: _Primitive
copy_p: _Primitive
cos_p: _Primitive
cosh_p: _Primitive
create_token_p: _Primitive
cumlogsumexp_p: _Primitive
cummax_p: _Primitive
cummin_p: _Primitive
cumprod_p: _Primitive
cumsum_p: _Primitive
device_put_p: _Primitive
digamma_p: _Primitive
div_p: _Primitive
dot_general_p: _Primitive
dynamic_slice_p: _Primitive
dynamic_update_slice_p: _Primitive
eq_p: _Primitive
eq_to_p: _Primitive
erf_inv_p: _Primitive
erf_p: _Primitive
erfc_p: _Primitive
exp2_p: _Primitive
exp_p: _Primitive
expm1_p: _Primitive
fft_p: _Primitive
floor_p: _Primitive
gather_p: _Primitive
ge_p: _Primitive
gt_p: _Primitive
igamma_grad_a_p: _Primitive
igamma_p: _Primitive
igammac_p: _Primitive
imag_p: _Primitive
integer_pow_p: _Primitive
iota_p: _Primitive
is_finite_p: _Primitive
le_p: _Primitive
le_to_p: _Primitive
lgamma_p: _Primitive
linear_solve_p: _Primitive
log1p_p: _Primitive
log2_p: _Primitive
log_p: _Primitive
logistic_p: _Primitive
lt_p: _Primitive
lt_to_p: _Primitive
max_p: _Primitive
min_p: _Primitive
mul_p: _Primitive
mulhi_p: _Primitive
ne_p: _Primitive
neg_p: _Primitive
nextafter_p: _Primitive
not_p: _Primitive
one_minus_square_p: _Primitive
optimization_barrier_p: _Primitive
or_p: _Primitive
pad_p: _Primitive
pmax_p: _Primitive
pmin_p: _Primitive
polygamma_p: _Primitive
polynomial_p: _Primitive
population_count_p: _Primitive
pow_p: _Primitive
ppermute_p: _Primitive
psum_p: _Primitive
ragged_all_to_all_p: _Primitive
real_p: _Primitive
reduce_and_p: _Primitive
reduce_max_p: _Primitive
reduce_min_p: _Primitive
reduce_or_p: _Primitive
reduce_p: _Primitive
reduce_precision_p: _Primitive
reduce_prod_p: _Primitive
reduce_sum_p: _Primitive
reduce_window_max_p: _Primitive
reduce_window_min_p: _Primitive
reduce_window_p: _Primitive
reduce_window_sum_p: _Primitive
reduce_xor_p: _Primitive
regularized_incomplete_beta_p: _Primitive
rem_p: _Primitive
reshape_p: _Primitive
rev_p: _Primitive
rng_bit_generator_p: _Primitive
rng_uniform_p: _Primitive
round_p: _Primitive
rsqrt_p: _Primitive
scan_p: _Primitive
scatter_add_p: _Primitive
scatter_max_p: _Primitive
scatter_min_p: _Primitive
scatter_mul_p: _Primitive
scatter_p: _Primitive
scatter_sub_p: _Primitive
select_and_gather_add_p: _Primitive
select_and_scatter_add_p: _Primitive
select_and_scatter_p: _Primitive
select_n_p: _Primitive
sharding_constraint_p: _Primitive
shift_left_p: _Primitive
shift_right_arithmetic_p: _Primitive
shift_right_logical_p: _Primitive
sign_p: _Primitive
sin_p: _Primitive
sinh_p: _Primitive
slice_p: _Primitive
sort_p: _Primitive
split_p: _Primitive
sqrt_p: _Primitive
square_p: _Primitive
squeeze_p: _Primitive
stack_p: _Primitive
stage_p: _Primitive
stop_gradient_p: _Primitive
sub_p: _Primitive
tan_p: _Primitive
tanh_p: _Primitive
tile_p: _Primitive
top_k_p: _Primitive
transpose_p: _Primitive
unstack_p: _Primitive
while_p: _Primitive
xor_p: _Primitive
zeta_p: _Primitive
3 changes: 3 additions & 0 deletions tensor-shapes/jax-pyrefly-stubs/test/test_lax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
import jax.lax as lax
import jax.numpy as jnp
import numpy as np
from jax.extend.core import Primitive
from shape_extensions import assert_raises, assert_shape, IntTuple, IntVar


Expand Down Expand Up @@ -1393,6 +1394,8 @@ def test_compiler_and_misc() -> None:
x, x, default=lambda a, b: a + b, cpu=lambda a, b: a - b
)
assert_shape(pd2.shape, (2, 3))
assert_type(lax.abs_p, Primitive)
assert isinstance(lax.abs_p, Primitive)


def test_lax_linalg_arraylike() -> None:
Expand Down
Loading