Coverage
October 25, 2021 · View on GitHub
This page documents which functions we provide shape-aware stubs for.
TensorFlow
Note that for TensorFlow, although we don't provide shape-aware stubs for functions marked ❌, we do provide normal stubs, so autocomplete etc should still work.
tensorflow
| a-g | g-r | r-z |
|---|---|---|
| ✅ abs | ❌ gather_nd | ✅ round |
| ✅ acos | ❌ get_static_value | ❌ saturate_cast |
| ✅ acosh | ❌ grad_pass_through | ❌ scalar_mul |
| ❌ add | ❌ gradients | ❌ scan |
| ❌ add_n | ❌ greater | ❌ scatter_nd |
| ❌ argmax | ❌ greater_equal | ❌ searchsorted |
| ❌ argmin | ❌ group | ❌ sequence_mask |
| ❌ argsort | ❌ hessians | ❌ shape |
| ❌ as_dtype | ❌ histogram_fixed_width | ❌ shape_n |
| ❌ as_string | ❌ histogram_fixed_width_bins | ✅ sigmoid |
| ✅ asin | ❌ identity | ✅ sign |
| ✅ asinh | ❌ identity_n | ✅ sin |
| ✅ atan | ❌ less | ✅ sinh |
| ❌ atan2 | ❌ less_equal | ❌ size |
| ✅ atanh | ❌ linspace | ❌ slice |
| ❌ batch_to_space | ❌ logical_and | ❌ sort |
| ❌ bitcast | ✅ logical_not | ❌ space_to_batch |
| ❌ boolean_mask | ❌ logical_or | ❌ space_to_batch_nd |
| ❌ broadcast_dynamic_shape | ❌ make_ndarray | ❌ split |
| ❌ broadcast_static_shape | ❌ map_fn | ✅ sqrt |
| ❌ broadcast_to | ✅ matmul | ✅ square |
| ❌ case | ❌ matrix_square_root | ❌ squeeze |
| ❌ cast | ❌ maximum | ❌ stack |
| ❌ clip_by_global_norm | ❌ meshgrid | ❌ stop_gradient |
| ❌ clip_by_norm | ❌ minimum | ❌ strided_slice |
| ❌ clip_by_value | ❌ multiply | ❌ subtract |
| ❌ complex | ✅ negative | ❌ switch_case |
| ❌ concat | ❌ norm | ✅ tan |
| ❌ cond | ❌ not_equal | ✅ tanh |
| ❌ constant | ❌ one_hot | ❌ tensor_scatter_nd_add |
| ❌ convert_to_tensor | ✅ ones | ❌ tensor_scatter_nd_max |
| ✅ cos | ✅ ones_like | ❌ tensor_scatter_nd_min |
| ✅ cosh | ❌ pad | ❌ tensor_scatter_nd_sub |
| ❌ cumsum | ❌ parallel_stack | ❌ tensor_scatter_nd_update |
| ❌ divide | ❌ pow | ❌ tensordot |
| ❌ dynamic_partition | ❌ range | ❌ tile |
| ❌ dynamic_stitch | ❌ rank | ❌ timestamp |
| ❌ edit_distance | ❌ realdiv | ✅ transpose |
| ❌ eig | ✅ reduce_all | ❌ truediv |
| ❌ eigvals | ✅ reduce_any | ❌ truncatediv |
| ❌ einsum | ✅ reduce_logsumexp | ❌ truncatemod |
| ❌ equal | ✅ reduce_max | ❌ tuple |
| ✅ exp | ✅ reduce_mean | ❌ unique |
| ❌ expand_dims | ✅ reduce_min | ❌ unique_with_counts |
| ❌ extract_volume_patches | ✅ reduce_prod | ❌ unravel_index |
| ❌ eye | ✅ reduce_sum | ❌ unstack |
| ❌ fill | ❌ repeat | ❌ vectorized_map |
| ❌ fingerprint | ❌ required_space_to_batch_paddings | ❌ where |
| ✅ floor | ❌ reshape | ❌ while_loop |
| ❌ foldl | ❌ reverse | ✅ zeros |
| ❌ foldr | ❌ reverse_sequence | ✅ zeros_like |
| ❌ gather | ❌ roll |
tensorflow.math
| a-g | g-r | r-z |
|---|---|---|
| ✅ abs | ❌ greater_equal | ❌ reduce_max |
| ❌ accumulate_n | ❌ igamma | ❌ reduce_mean |
| ✅ acos | ❌ igammac | ❌ reduce_min |
| ✅ acosh | ❌ imag | ❌ reduce_prod |
| ❌ add | ❌ in_top_k | ❌ reduce_std |
| ❌ add_n | ❌ invert_permutation | ❌ reduce_sum |
| ❌ angle | ✅ is_finite | ❌ reduce_variance |
| ❌ argmax | ✅ is_inf | ✅ rint |
| ❌ argmin | ✅ is_nan | ✅ round |
| ✅ asin | ❌ is_non_decreasing | ✅ rsqrt |
| ✅ asinh | ❌ is_strictly_increasing | ❌ scalar_mul |
| ✅ atan | ❌ l2_normalize | ❌ segment_max |
| ❌ atan2 | ✅ lbeta | ❌ segment_mean |
| ✅ atanh | ❌ less | ❌ segment_min |
| ❌ bessel_i0 | ❌ less_equal | ❌ segment_prod |
| ❌ bessel_i0e | ✅ lgamma | ❌ segment_sum |
| ❌ bessel_i1 | ✅ log | ✅ sigmoid |
| ❌ bessel_i1e | ✅ log1p | ✅ sign |
| ❌ betainc | ✅ log_sigmoid | ✅ sin |
| ❌ bincount | ❌ log_softmax | ✅ sinh |
| ❌ ceil | ❌ logical_and | ❌ sobol_sample |
| ❌ confusion_matrix | ✅ logical_not | ❌ softmax |
| ❌ conj | ❌ logical_or | ✅ softplus |
| ✅ cos | ❌ logical_xor | ✅ softsign |
| ✅ cosh | ❌ maximum | ✅ sqrt |
| ❌ count_nonzero | ❌ minimum | ✅ square |
| ❌ cumprod | ❌ mod | ❌ squared_difference |
| ❌ cumsum | ❌ multiply | ❌ subtract |
| ❌ cumulative_logsumexp | ❌ multiply_no_nan | ✅ tan |
| ❌ digamma | ✅ ndtri | ✅ tanh |
| ❌ divide | ✅ negative | ❌ top_k |
| ❌ divide_no_nan | ❌ nextafter | ❌ truediv |
| ❌ equal | ❌ not_equal | ❌ unsorted_segment_max |
| ✅ erf | ❌ polygamma | ❌ unsorted_segment_mean |
| ✅ erfc | ❌ polyval | ❌ unsorted_segment_min |
| ❌ erfcinv | ❌ pow | ❌ unsorted_segment_prod |
| ✅ erfinv | ❌ real | ❌ unsorted_segment_sqrt_n |
| ✅ exp | ✅ reciprocal | ❌ unsorted_segment_sum |
| ✅ expm1 | ✅ reciprocal_no_nan | ❌ xdivy |
| ✅ floor | ❌ reduce_all | ❌ xlog1py |
| ❌ floordiv | ❌ reduce_any | ❌ xlogy |
| ❌ floormod | ❌ reduce_euclidean_norm | ❌ zero_fraction |
| ❌ greater | ❌ reduce_logsumexp | ❌ zeta |
JAX
jax.numpy
| a-f | f-n | n-z |
|---|---|---|
| ✅ abs | ❌ floor_divide | ❌ ndfromtxt |
| ❌ absolute | ❌ fmax | ❌ ndim |
| ❌ add | ❌ fmin | ✅ negative |
| ❌ all | ❌ fmod | ❌ nested_iters |
| ❌ allclose | ❌ format_float_positional | ❌ nextafter |
| ❌ alltrue | ❌ format_float_scientific | ❌ nonzero |
| ❌ amax | ❌ frexp | ❌ not_equal |
| ❌ amin | ❌ frombuffer | ❌ nper |
| ❌ angle | ❌ fromfile | ❌ npv |
| ❌ any | ❌ fromiter | ❌ ones |
| ❌ append | ❌ frompyfunc | ✅ ones_like |
| ❌ apply_along_axis | ❌ fromregex | ❌ outer |
| ❌ apply_over_axes | ❌ fromstring | ❌ packbits |
| ❌ arange | ❌ full | ❌ pad |
| ❌ arccos | ❌ full_like | ❌ partition |
| ❌ arccosh | ❌ fv | ❌ percentile |
| ❌ arcsin | ❌ gcd | ❌ piecewise |
| ❌ arcsinh | ❌ genfromtxt | ❌ place |
| ❌ arctan | ❌ geomspace | ❌ pmt |
| ❌ arctan2 | ❌ get_array_wrap | ❌ poly |
| ❌ arctanh | ❌ get_include | ❌ polyadd |
| ❌ argmax | ❌ get_printoptions | ❌ polyder |
| ❌ argmin | ❌ getbufsize | ❌ polydiv |
| ❌ argpartition | ❌ geterr | ❌ polyfit |
| ❌ argsort | ❌ geterrcall | ❌ polyint |
| ❌ argwhere | ❌ geterrobj | ❌ polymul |
| ❌ around | ❌ gradient | ❌ polysub |
| ❌ array | ❌ greater | ❌ polyval |
| ❌ array2string | ❌ greater_equal | ❌ positive |
| ❌ array_equal | ❌ hamming | ❌ power |
| ❌ array_equiv | ❌ hanning | ❌ ppmt |
| ❌ array_repr | ❌ heaviside | ❌ printoptions |
| ❌ array_split | ❌ histogram | ❌ prod |
| ❌ array_str | ❌ histogram2d | ❌ product |
| ❌ asanyarray | ❌ histogram_bin_edges | ❌ promote_types |
| ❌ asarray | ❌ histogramdd | ❌ ptp |
| ❌ asarray_chkfinite | ❌ hsplit | ❌ put |
| ❌ ascontiguousarray | ❌ hstack | ❌ put_along_axis |
| ❌ asfarray | ❌ hypot | ❌ putmask |
| ❌ asfortranarray | ❌ i0 | ❌ pv |
| ❌ asmatrix | ❌ identity | ❌ quantile |
| ❌ asscalar | ❌ imag | ❌ rad2deg |
| ❌ atleast_1d | ❌ in1d | ❌ radians |
| ❌ atleast_2d | ❌ indices | ❌ rate |
| ❌ atleast_3d | ❌ info | ❌ ravel |
| ❌ average | ❌ inner | ❌ ravel_multi_index |
| ❌ bartlett | ❌ insert | ❌ real |
| ❌ base_repr | ❌ interp | ❌ real_if_close |
| ❌ binary_repr | ❌ intersect1d | ❌ recfromcsv |
| ❌ bincount | ❌ invert | ❌ recfromtxt |
| ❌ bitwise_and | ❌ ipmt | ❌ reciprocal |
| ❌ bitwise_not | ❌ irr | ❌ remainder |
| ❌ bitwise_or | ❌ is_busday | ❌ repeat |
| ❌ bitwise_xor | ❌ isclose | ❌ require |
| ❌ blackman | ❌ iscomplex | ❌ reshape |
| ❌ block | ❌ iscomplexobj | ❌ resize |
| ❌ bmat | ❌ isfinite | ❌ result_type |
| ❌ broadcast_arrays | ❌ isfortran | ❌ right_shift |
| ❌ broadcast_shapes | ❌ isin | ❌ rint |
| ❌ broadcast_to | ❌ isinf | ❌ roll |
| ❌ busday_count | ❌ isnan | ❌ rollaxis |
| ❌ busday_offset | ❌ isnat | ❌ roots |
| ❌ byte_bounds | ❌ isneginf | ❌ rot90 |
| ❌ can_cast | ❌ isposinf | ✅ round |
| ❌ cbrt | ❌ isreal | ❌ round_ |
| ❌ ceil | ❌ isrealobj | ❌ row_stack |
| ❌ choose | ❌ isscalar | ❌ safe_eval |
| ❌ clip | ❌ issctype | ❌ savez |
| ❌ column_stack | ❌ issubclass_ | ❌ searchsorted |
| ❌ common_type | ❌ issubdtype | ❌ select |
| ❌ compare_chararrays | ❌ issubsctype | ❌ setbufsize |
| ❌ compress | ❌ iterable | ❌ setdiff1d |
| ❌ concatenate | ❌ ix_ | ❌ setxor1d |
| ❌ conj | ❌ kaiser | ❌ shape |
| ❌ conjugate | ❌ kron | ✅ sign |
| ❌ convolve | ❌ lcm | ❌ signbit |
| ❌ copy | ❌ ldexp | ✅ sin |
| ❌ copysign | ❌ left_shift | ❌ sinc |
| ❌ copyto | ❌ less | ✅ sinh |
| ❌ corrcoef | ❌ less_equal | ❌ size |
| ❌ correlate | ❌ lexsort | ❌ sometrue |
| ✅ cos | ❌ linspace | ❌ sort |
| ✅ cosh | ❌ load | ❌ sort_complex |
| ❌ count_nonzero | ❌ loads | ❌ source |
| ❌ cov | ❌ loadtxt | ❌ spacing |
| ❌ cross | ❌ log | ❌ split |
| ❌ cumprod | ❌ log10 | ✅ sqrt |
| ❌ cumproduct | ❌ log1p | ✅ square |
| ❌ cumsum | ❌ log2 | ❌ squeeze |
| ❌ datetime_as_string | ❌ logaddexp | ❌ stack |
| ❌ datetime_data | ❌ logaddexp2 | ❌ std |
| ❌ deg2rad | ❌ logical_and | ❌ subtract |
| ❌ degrees | ✅ logical_not | ❌ sum |
| ❌ delete | ❌ logical_or | ❌ swapaxes |
| ❌ diag | ❌ logical_xor | ❌ take |
| ❌ diag_indices | ❌ logspace | ❌ take_along_axis |
| ❌ diag_indices_from | ❌ lookfor | ✅ tan |
| ❌ diagflat | ❌ mafromtxt | ✅ tanh |
| ❌ diagonal | ❌ mask_indices | ❌ tensordot |
| ❌ diff | ❌ mat | ❌ tile |
| ❌ digitize | ❌ matmul | ❌ trace |
| ❌ disp | ❌ max | ✅ transpose |
| ❌ divide | ❌ maximum | ❌ trapz |
| ❌ divmod | ❌ mean | ❌ tri |
| ❌ dot | ❌ median | ❌ tril |
| ❌ dsplit | ❌ meshgrid | ❌ tril_indices |
| ❌ dstack | ❌ min | ❌ tril_indices_from |
| ❌ ediff1d | ❌ min_scalar_type | ❌ trim_zeros |
| ❌ einsum | ❌ minimum | ❌ triu |
| ❌ einsum_path | ❌ mirr | ❌ triu_indices |
| ❌ empty | ❌ mod | ❌ triu_indices_from |
| ❌ empty_like | ❌ modf | ❌ true_divide |
| ❌ equal | ❌ moveaxis | ❌ trunc |
| ✅ exp | ❌ msort | ❌ typename |
| ❌ exp2 | ❌ multiply | ❌ union1d |
| ❌ expand_dims | ❌ nan_to_num | ❌ unique |
| ❌ expm1 | ❌ nanargmax | ❌ unpackbits |
| ❌ extract | ❌ nanargmin | ❌ unravel_index |
| ❌ eye | ❌ nancumprod | ❌ unwrap |
| ❌ fabs | ❌ nancumsum | ❌ vander |
| ❌ fastCopyAndTranspose | ❌ nanmax | ❌ var |
| ❌ fill_diagonal | ❌ nanmean | ❌ vdot |
| ❌ find_common_type | ❌ nanmedian | ❌ vectorize |
| ❌ fix | ❌ nanmin | ❌ vsplit |
| ❌ flatnonzero | ❌ nanpercentile | ❌ vstack |
| ❌ flip | ❌ nanprod | ❌ where |
| ❌ fliplr | ❌ nanquantile | ❌ who |
| ❌ flipud | ❌ nanstd | ❌ zeros |
| ❌ float_power | ❌ nansum | ✅ zeros_like |
| ✅ floor | ❌ nanvar |