{ "package": "jax", "version": "0.9.1", "stubs_only": "no", "py_typed": "YES", "pypi": { "upload_time": "2026-03-02T11:24:18.382976Z", "requires_python": ">=3.11", "size": 2625994, "sha256": "ce1b82477ee192f0b1d9801b095aa0cf3839bc1fe0cbc071c961a24b3ff30361" }, "metadata": { "Metadata-Version": [ "2.4" ], "Name": [ "jax" ], "Version": [ "0.9.1" ], "Summary": [ "Differentiate, compile, and transform Numpy code." ], "Home-page": [ "https://github.com/jax-ml/jax" ], "Author": [ "JAX team" ], "Author-email": [ "jax-dev@google.com" ], "License": [ "Apache-2.0" ], "Classifier": [ "Development Status :: 5 - Production/Stable", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.14", "Programming Language :: Python :: Free Threading :: 3 - Stable" ], "Requires-Python": [ ">=3.11" ], "Description-Content-Type": [ "text/markdown" ], "License-File": [ "LICENSE", "AUTHORS" ], "Requires-Dist": [ "jaxlib<=0.9.1,>=0.9.1", "ml_dtypes>=0.5.0", "numpy>=2.0", "opt_einsum", "scipy>=1.13", "jaxlib==0.9.1; extra == \"minimum-jaxlib\"", "jaxlib==0.9.0; extra == \"ci\"", "jaxlib<=0.9.1,>=0.9.1; extra == \"tpu\"", "libtpu==0.0.35.*; extra == \"tpu\"", "requests; extra == \"tpu\"", "jaxlib<=0.9.1,>=0.9.1; extra == \"cuda\"", "jax-cuda12-plugin[with-cuda]<=0.9.1,>=0.9.1; extra == \"cuda\"", "jaxlib<=0.9.1,>=0.9.1; extra == \"cuda12\"", "jax-cuda12-plugin[with-cuda]<=0.9.1,>=0.9.1; extra == \"cuda12\"", "jaxlib<=0.9.1,>=0.9.1; extra == \"cuda13\"", "jax-cuda13-plugin[with-cuda]<=0.9.1,>=0.9.1; extra == \"cuda13\"", "jaxlib<=0.9.1,>=0.9.1; extra == \"cuda12-local\"", "jax-cuda12-plugin<=0.9.1,>=0.9.1; extra == \"cuda12-local\"", "jaxlib<=0.9.1,>=0.9.1; extra == \"cuda13-local\"", "jax-cuda13-plugin<=0.9.1,>=0.9.1; extra == \"cuda13-local\"", "jaxlib<=0.9.1,>=0.9.1; extra == \"rocm\"", "jax-rocm7-plugin<=0.9.1,>=0.9.1; extra == \"rocm\"", "kubernetes; extra == \"k8s\"", "xprof; extra == \"xprof\"" ], "Provides-Extra": [ "minimum-jaxlib", "cpu", "ci", "tpu", "cuda", "cuda12", "cuda13", "cuda12-local", "cuda13-local", "rocm", "k8s", "xprof" ], "Dynamic": [ "author", "author-email", "classifier", "description", "description-content-type", "home-page", "license", "license-file", "provides-extra", "requires-dist", "requires-python", "summary" ] }, "module_reports": [ { "path": "jax/_src/pjit.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pjit.reshard", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pjit.with_layout_constraint", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax._src.pjit.jit_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.pjit.sharding_constraint_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.pjit.with_sharding_constraint", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pjit.auto_axes", "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pjit.explicit_axes", "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [ { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null } ], "name": "jax._src.pjit", "names": [ "jax._src.pjit.auto_axes", "jax._src.pjit.explicit_axes", "jax._src.pjit.jit_p", "jax._src.pjit.reshard", "jax._src.pjit.sharding_constraint_p", "jax._src.pjit.with_layout_constraint", "jax._src.pjit.with_sharding_constraint" ], "n_typable": 19, "n_typed": 2, "n_any": 0, "n_untyped": 17, "n_functions": 5, "n_function_overloads": 5, "n_function_params": 12, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 2, "n_properties": 0, "n_type_ignores": 5 }, { "path": "jax/_src/api.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.api.vjp", "n_typed": 4, "n_any": 1, "n_untyped": 0, "n_overloads": 2, "n_typable": 5, "n_params": 4, "n_function_overloads": 2, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.api.jacrev", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.api.jacfwd", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.api.clear_caches", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.api.vmap", "n_typed": 7, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.api.jit", "n_typed": 11, "n_any": 2, "n_untyped": 0, "n_overloads": 2, "n_typable": 13, "n_params": 12, "n_function_overloads": 2, "n_function_params": 12 }, { "kind": "function", "name": "jax._src.api.named_scope", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.api.disable_jit", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.api.hessian", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.api.make_jaxpr", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 2, "n_typable": 5, "n_params": 4, "n_function_overloads": 2, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.api.eval_shape", "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.api.effects_barrier", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.api.pmap", "n_typed": 9, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 11, "n_params": 10, "n_function_overloads": 1, "n_function_params": 10 }, { "kind": "function", "name": "jax._src.api.device_put", "n_typed": 4, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.api.device_get", "n_typed": 0, "n_any": 1, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.api.jvp", "n_typed": 3, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.api.jacobian", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.api.value_and_grad", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.api.linearize", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 2, "n_typable": 4, "n_params": 3, "n_function_overloads": 2, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.api.grad", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.api.linear_transpose", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.api.block_until_ready", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.api.copy_to_host_async", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.api.named_call", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.api.fwd_and_bwd", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.api.live_arrays", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.api.clear_backends", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "type_ignores": [ { "kind": "type", "rules": null }, { "kind": "type", "rules": null } ], "name": "jax._src.api", "names": [ "jax._src.api.block_until_ready", "jax._src.api.clear_backends", "jax._src.api.clear_caches", "jax._src.api.copy_to_host_async", "jax._src.api.device_get", "jax._src.api.device_put", "jax._src.api.disable_jit", "jax._src.api.effects_barrier", "jax._src.api.eval_shape", "jax._src.api.fwd_and_bwd", "jax._src.api.grad", "jax._src.api.hessian", "jax._src.api.jacfwd", "jax._src.api.jacobian", "jax._src.api.jacrev", "jax._src.api.jit", "jax._src.api.jvp", "jax._src.api.linear_transpose", "jax._src.api.linearize", "jax._src.api.live_arrays", "jax._src.api.make_jaxpr", "jax._src.api.named_call", "jax._src.api.named_scope", "jax._src.api.pmap", "jax._src.api.value_and_grad", "jax._src.api.vjp", "jax._src.api.vmap" ], "n_typable": 124, "n_typed": 96, "n_any": 7, "n_untyped": 21, "n_functions": 27, "n_function_overloads": 31, "n_function_params": 97, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 2 }, { "path": "jax/_src/array.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.array.make_array_from_callback", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.array.make_array_from_process_local_data", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax._src.array.Shard", "methods": [ { "kind": "function", "name": "Shard.__init__", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "Shard.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "Shard.index", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Shard.replica_id", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Shard.device", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Shard.data", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [], "n_typable": 10, "n_typed": 6, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 4, "n_attrs": 0, "n_properties": 4 }, { "kind": "function", "name": "jax._src.array.make_array_from_single_device_arrays", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [ { "kind": "type", "rules": null } ], "name": "jax._src.array", "names": [ "jax._src.array.Shard", "jax._src.array.make_array_from_callback", "jax._src.array.make_array_from_process_local_data", "jax._src.array.make_array_from_single_device_arrays" ], "n_typable": 24, "n_typed": 16, "n_any": 0, "n_untyped": 8, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 11, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 4, "n_classes": 1, "n_attrs": 0, "n_properties": 4, "n_type_ignores": 1 }, { "path": "jax/_src/xla_bridge.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.xla_bridge.local_device_count", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.xla_bridge.devices", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.xla_bridge.device_count", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.xla_bridge.default_backend", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.xla_bridge.process_index", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.xla_bridge.host_ids", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.xla_bridge.host_count", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.xla_bridge.process_indices", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.xla_bridge.process_count", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.xla_bridge.local_devices", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.xla_bridge.host_id", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.xla_bridge.backend_xla_version", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.xla_bridge.get_backend", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.xla_bridge.register_backend_factory", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.xla_bridge.backends", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "type_ignores": [], "name": "jax._src.xla_bridge", "names": [ "jax._src.xla_bridge.backend_xla_version", "jax._src.xla_bridge.backends", "jax._src.xla_bridge.default_backend", "jax._src.xla_bridge.device_count", "jax._src.xla_bridge.devices", "jax._src.xla_bridge.get_backend", "jax._src.xla_bridge.host_count", "jax._src.xla_bridge.host_id", "jax._src.xla_bridge.host_ids", "jax._src.xla_bridge.local_device_count", "jax._src.xla_bridge.local_devices", "jax._src.xla_bridge.process_count", "jax._src.xla_bridge.process_index", "jax._src.xla_bridge.process_indices", "jax._src.xla_bridge.register_backend_factory" ], "n_typable": 36, "n_typed": 35, "n_any": 0, "n_untyped": 1, "n_functions": 15, "n_function_overloads": 15, "n_function_params": 21, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/sharding_impls.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.sharding_impls.set_mesh", "methods": [ { "kind": "function", "name": "set_mesh.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "set_mesh.__enter__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "set_mesh.__exit__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "set_mesh.prev_abstract_mesh", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "set_mesh.prev_mesh", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 9, "n_typed": 1, "n_any": 0, "n_untyped": 8, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 4, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax._src.sharding_impls.make_mesh", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "class", "name": "jax._src.sharding_impls.GSPMDSharding", "methods": [ { "kind": "function", "name": "GSPMDSharding.__init__", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "GSPMDSharding.__reduce__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "GSPMDSharding.__eq__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "GSPMDSharding.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "GSPMDSharding.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "GSPMDSharding.check_compatible_aval", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "GSPMDSharding.with_memory_kind", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "GSPMDSharding._to_xla_hlo_sharding", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "GSPMDSharding._to_sdy_sharding", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "GSPMDSharding.get_replicated", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [ { "kind": "property", "name": "GSPMDSharding._hlo_sharding_hash", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "GSPMDSharding.num_devices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "GSPMDSharding.device_set", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "GSPMDSharding.memory_kind", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "GSPMDSharding._device_assignment", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "GSPMDSharding.is_fully_replicated", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "GSPMDSharding.is_fully_addressable", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "GSPMDSharding._devices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "GSPMDSharding._hlo_sharding", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "GSPMDSharding._memory_kind", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "GSPMDSharding._internal_device_list", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 31, "n_typed": 22, "n_any": 0, "n_untyped": 9, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 10, "n_method_overloads": 10, "n_method_params": 10, "n_attrs": 4, "n_properties": 7 }, { "kind": "class", "name": "jax._src.sharding_impls.SPMDAxisContext", "methods": [ { "kind": "function", "name": "SPMDAxisContext.extend_manual", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "SPMDAxisContext.axis_env", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "SPMDAxisContext.unsafe_axis_env", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "SPMDAxisContext.mesh", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SPMDAxisContext.manual_axes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 4, "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 1, "n_attrs": 2, "n_properties": 2 }, { "kind": "class", "name": "jax._src.sharding_impls.ShardingContext", "methods": [ { "kind": "function", "name": "ShardingContext.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "ShardingContext.axis_env", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "ShardingContext.num_devices", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ShardingContext.device_assignment", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ShardingContext.abstract_mesh", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 2, "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 3, "n_properties": 1 }, { "kind": "class", "name": "jax._src.sharding_impls.ReplicaAxisContext", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "ReplicaAxisContext.axis_env", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 1, "n_properties": 0 }, { "kind": "class", "name": "jax._src.sharding_impls.SingleDeviceSharding", "methods": [ { "kind": "function", "name": "SingleDeviceSharding.__init__", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "SingleDeviceSharding.__reduce__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "SingleDeviceSharding.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "SingleDeviceSharding.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "SingleDeviceSharding.__eq__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SingleDeviceSharding.with_memory_kind", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SingleDeviceSharding.devices_indices_map", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SingleDeviceSharding._to_xla_hlo_sharding", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SingleDeviceSharding._to_sdy_sharding", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SingleDeviceSharding.check_compatible_aval", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "SingleDeviceSharding.num_devices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "SingleDeviceSharding.device_set", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "SingleDeviceSharding.memory_kind", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "SingleDeviceSharding._device_assignment", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "SingleDeviceSharding.is_fully_replicated", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "SingleDeviceSharding.is_fully_addressable", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "SingleDeviceSharding._device", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "SingleDeviceSharding._memory_kind", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 26, "n_typed": 20, "n_any": 0, "n_untyped": 6, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 10, "n_method_overloads": 10, "n_method_params": 8, "n_attrs": 2, "n_properties": 6 }, { "kind": "function", "name": "jax._src.sharding_impls.get_mesh", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "type_ignores": [ { "kind": "type", "rules": null } ], "name": "jax._src.sharding_impls", "names": [ "jax._src.sharding_impls.GSPMDSharding", "jax._src.sharding_impls.ReplicaAxisContext", "jax._src.sharding_impls.SPMDAxisContext", "jax._src.sharding_impls.ShardingContext", "jax._src.sharding_impls.SingleDeviceSharding", "jax._src.sharding_impls.get_mesh", "jax._src.sharding_impls.make_mesh", "jax._src.sharding_impls.set_mesh" ], "n_typable": 78, "n_typed": 51, "n_any": 0, "n_untyped": 27, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 4, "n_methods": 25, "n_method_overloads": 25, "n_method_params": 23, "n_classes": 6, "n_attrs": 14, "n_properties": 16, "n_type_ignores": 1 }, { "path": "jax/_src/config.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.config.check_tracer_leaks", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.threefry_partitionable", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.checking_leaks", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.debug_infs", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.array_garbage_collection_guard", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.numpy_dtype_promotion", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.transfer_guard_device_to_device", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.remove_size_one_mesh_axis_from_type", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.enable_custom_prng", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.thread_guard", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.numpy_rank_promotion", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.no_tracing", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.no_execution", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.config.make_user_context", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax._src.config.softmax_custom_jvp", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.default_matmul_precision", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.transfer_guard_device_to_host", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.default_prng_impl", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.legacy_prng_key", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.default_device", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.debug_nans", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.debug_key_reuse", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.jax2tf_associative_scan_reductions", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.allow_f16_reductions", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.transfer_guard_host_to_device", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.config", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.enable_x64", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.enable_checks", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.explain_cache_misses", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.config.log_compiles", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.config.transfer_guard", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax._src.config.enable_custom_vjp_by_custom_transpose", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [ { "kind": "type", "rules": null } ], "name": "jax._src.config", "names": [ "jax._src.config.allow_f16_reductions", "jax._src.config.array_garbage_collection_guard", "jax._src.config.check_tracer_leaks", "jax._src.config.checking_leaks", "jax._src.config.config", "jax._src.config.debug_infs", "jax._src.config.debug_key_reuse", "jax._src.config.debug_nans", "jax._src.config.default_device", "jax._src.config.default_matmul_precision", "jax._src.config.default_prng_impl", "jax._src.config.enable_checks", "jax._src.config.enable_custom_prng", "jax._src.config.enable_custom_vjp_by_custom_transpose", "jax._src.config.enable_x64", "jax._src.config.explain_cache_misses", "jax._src.config.jax2tf_associative_scan_reductions", "jax._src.config.legacy_prng_key", "jax._src.config.log_compiles", "jax._src.config.make_user_context", "jax._src.config.no_execution", "jax._src.config.no_tracing", "jax._src.config.numpy_dtype_promotion", "jax._src.config.numpy_rank_promotion", "jax._src.config.remove_size_one_mesh_axis_from_type", "jax._src.config.softmax_custom_jvp", "jax._src.config.thread_guard", "jax._src.config.threefry_partitionable", "jax._src.config.transfer_guard", "jax._src.config.transfer_guard_device_to_device", "jax._src.config.transfer_guard_device_to_host", "jax._src.config.transfer_guard_host_to_device" ], "n_typable": 34, "n_typed": 2, "n_any": 0, "n_untyped": 32, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 2, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 30, "n_properties": 0, "n_type_ignores": 1 }, { "path": "jax/_src/custom_derivatives.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.custom_derivatives.custom_gradient", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax._src.custom_derivatives.custom_jvp", "methods": [ { "kind": "function", "name": "custom_jvp.__init__", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "custom_jvp.defjvp", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "custom_jvp.defjvps", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "custom_jvp.__call__", "n_typed": 1, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "custom_jvp.fun", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "custom_jvp.nondiff_argnums", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "custom_jvp.nondiff_argnames", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "custom_jvp.jvp", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "custom_jvp.symbolic_zeros", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "custom_jvp.__getattr__", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 17, "n_typed": 14, "n_any": 2, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 8, "n_attrs": 6, "n_properties": 0 }, { "kind": "class", "name": "jax._src.custom_derivatives.custom_vjp", "methods": [ { "kind": "function", "name": "custom_vjp.__new__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "custom_vjp.__init__", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "custom_vjp.defvjp", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "custom_vjp.__call__", "n_typed": 1, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "custom_vjp.fun", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "custom_vjp.nondiff_argnums", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "custom_vjp.fwd", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "custom_vjp.bwd", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "custom_vjp.symbolic_zeros", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "custom_vjp.optimize_remat", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "custom_vjp.__getattr__", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 22, "n_typed": 11, "n_any": 2, "n_untyped": 9, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 12, "n_attrs": 7, "n_properties": 0 }, { "kind": "function", "name": "jax._src.custom_derivatives.closure_convert", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.custom_derivatives.CustomVJPPrimal", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "CustomVJPPrimal.value", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CustomVJPPrimal.perturbed", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "attr", "name": "jax._src.custom_derivatives.remat_opt_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.custom_derivatives.custom_jvp_call_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.custom_derivatives.custom_vjp_call_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.custom_derivatives.linear_call", "n_typed": 2, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.custom_derivatives.custom_vjp_primal_tree_values", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax._src.custom_derivatives", "names": [ "jax._src.custom_derivatives.CustomVJPPrimal", "jax._src.custom_derivatives.closure_convert", "jax._src.custom_derivatives.custom_gradient", "jax._src.custom_derivatives.custom_jvp", "jax._src.custom_derivatives.custom_jvp_call_p", "jax._src.custom_derivatives.custom_vjp", "jax._src.custom_derivatives.custom_vjp_call_p", "jax._src.custom_derivatives.custom_vjp_primal_tree_values", "jax._src.custom_derivatives.linear_call", "jax._src.custom_derivatives.remat_opt_p" ], "n_typable": 54, "n_typed": 29, "n_any": 4, "n_untyped": 21, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 8, "n_methods": 8, "n_method_overloads": 8, "n_method_params": 20, "n_classes": 3, "n_attrs": 18, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/callback.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.callback.pure_callback", "n_typed": 4, "n_any": 3, "n_untyped": 1, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.callback.io_callback", "n_typed": 3, "n_any": 3, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.callback.emit_python_callback", "n_typed": 10, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 11, "n_params": 10, "n_function_overloads": 1, "n_function_params": 10 } ], "type_ignores": [], "name": "jax._src.callback", "names": [ "jax._src.callback.emit_python_callback", "jax._src.callback.io_callback", "jax._src.callback.pure_callback" ], "n_typable": 26, "n_typed": 17, "n_any": 6, "n_untyped": 3, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 23, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/basearray.pyi", "symbol_reports": [ { "kind": "class", "name": "jax._src.basearray.Array", "methods": [ { "kind": "function", "name": "Array.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "Array.__array_namespace__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__getitem__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__setitem__", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Array.__len__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.__iter__", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.__reversed__", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.__round__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__lt__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__le__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__eq__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__ne__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__gt__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__ge__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__neg__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.__pos__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.__abs__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.__invert__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.__add__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__sub__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__mul__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__matmul__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__truediv__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__floordiv__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__mod__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__divmod__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__pow__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__lshift__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__rshift__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__and__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__xor__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__or__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__radd__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__rsub__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__rmul__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__rmatmul__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__rtruediv__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__rfloordiv__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__rmod__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__rdivmod__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__rpow__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__rlshift__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__rrshift__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__rand__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__rxor__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__ror__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__bool__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.__complex__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.__int__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.__float__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.__index__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.__buffer__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.__release_buffer__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.all", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "Array.any", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "Array.argmax", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Array.argmin", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Array.argpartition", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Array.argsort", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "Array.astype", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Array.choose", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Array.clip", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Array.compress", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "Array.conj", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.conjugate", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.copy", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.cumprod", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Array.cumsum", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Array.diagonal", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Array.dot", "n_typed": 3, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Array.flatten", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.item", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.max", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "Array.mean", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "Array.min", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "Array.nonzero", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Array.prod", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "Array.ptp", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Array.ravel", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Array.repeat", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "Array.reshape", "n_typed": 3, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Array.round", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Array.searchsorted", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "Array.sort", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "Array.squeeze", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.std", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "Array.sum", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "Array.swapaxes", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Array.take", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "Array.tobytes", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.tolist", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.trace", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "Array.transpose", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.var", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "Array.view", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Array.__array__", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Array.__dlpack__", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.addressable_data", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Array.block_until_ready", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.copy_to_host_async", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.delete", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array._rewrap_with_aval_and_sharding", "n_typed": 1, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Array.devices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.is_deleted", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.on_device_size_in_bytes", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.unsafe_buffer_pointer", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Array.to_device", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [ { "kind": "property", "name": "Array.dtype", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.__numpy_dtype__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.ndim", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.size", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.itemsize", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.imag", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.nbytes", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.real", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.T", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.mT", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.at", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.weak_type", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.sharding", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.committed", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.device", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.addressable_shards", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.global_shards", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.is_fully_addressable", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.is_fully_replicated", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Array.traceback", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "Array.aval", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 318, "n_typed": 296, "n_any": 11, "n_untyped": 11, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 107, "n_method_overloads": 107, "n_method_params": 189, "n_attrs": 1, "n_properties": 21 }, { "kind": "attr", "name": "jax._src.basearray.ArrayLike", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "type_ignores": [ { "kind": "type", "rules": [ "override" ] }, { "kind": "type", "rules": [ "override" ] }, { "kind": "type", "rules": [ "misc" ] }, { "kind": "type", "rules": [ "misc" ] }, { "kind": "type", "rules": [ "misc" ] }, { "kind": "type", "rules": [ "misc" ] }, { "kind": "type", "rules": [ "misc" ] } ], "name": "jax._src.basearray", "names": [ "jax._src.basearray.Array", "jax._src.basearray.ArrayLike" ], "n_typable": 319, "n_typed": 297, "n_any": 11, "n_untyped": 11, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 107, "n_method_overloads": 107, "n_method_params": 189, "n_classes": 1, "n_attrs": 2, "n_properties": 21, "n_type_ignores": 7 }, { "path": "jax/_src/dtypes.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.dtypes.float0", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "class", "name": "jax._src.dtypes.extended", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax._src.dtypes.scalar_type_of", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.dtypes.itemsize_bits", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.dtypes.canonicalize_dtype", "n_typed": 2, "n_any": 1, "n_untyped": 0, "n_overloads": 2, "n_typable": 3, "n_params": 2, "n_function_overloads": 2, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.dtypes.result_type", "n_typed": 2, "n_any": 1, "n_untyped": 0, "n_overloads": 3, "n_typable": 3, "n_params": 2, "n_function_overloads": 3, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.dtypes.prng_key", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "attr", "name": "jax._src.dtypes.bfloat16", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax._src.dtypes.primal_tangent_dtype", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax._src.dtypes.canonicalize_value_handlers", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "type_ignores": [ { "kind": "type", "rules": null } ], "name": "jax._src.dtypes", "names": [ "jax._src.dtypes.bfloat16", "jax._src.dtypes.canonicalize_dtype", "jax._src.dtypes.canonicalize_value_handlers", "jax._src.dtypes.extended", "jax._src.dtypes.float0", "jax._src.dtypes.itemsize_bits", "jax._src.dtypes.primal_tangent_dtype", "jax._src.dtypes.prng_key", "jax._src.dtypes.result_type", "jax._src.dtypes.scalar_type_of" ], "n_typable": 17, "n_typed": 12, "n_any": 3, "n_untyped": 2, "n_functions": 5, "n_function_overloads": 8, "n_function_params": 9, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 2, "n_attrs": 3, "n_properties": 0, "n_type_ignores": 1 }, { "path": "jax/_src/partition_spec.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.partition_spec.P", "methods": [ { "kind": "function", "name": "P.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "P.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "P.__reduce__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "P.__getitem__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "P.__iter__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "P.__len__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "P.__eq__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "P.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "P.__add__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "P.__radd__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "P.index", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "P.count", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "P.update", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "P.to_lo", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "P._normalized_spec_for_aval", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "P._check_compatible_wrt_shape", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "P.__match_args__", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "P.UNCONSTRAINED", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "P.unreduced", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "P.reduced", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 30, "n_typed": 2, "n_any": 0, "n_untyped": 28, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 16, "n_method_overloads": 16, "n_method_params": 12, "n_attrs": 4, "n_properties": 0 } ], "type_ignores": [], "name": "jax._src.partition_spec", "names": [ "jax._src.partition_spec.P" ], "n_typable": 30, "n_typed": 2, "n_any": 0, "n_untyped": 28, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 16, "n_method_overloads": 16, "n_method_params": 12, "n_classes": 1, "n_attrs": 4, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/core.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.core.freeze", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.core.empty_ref", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.core.ShapeDtypeStruct", "methods": [ { "kind": "function", "name": "ShapeDtypeStruct.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "ShapeDtypeStruct.__len__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ShapeDtypeStruct.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ShapeDtypeStruct.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ShapeDtypeStruct.__eq__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ShapeDtypeStruct.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ShapeDtypeStruct.__setattr__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "ShapeDtypeStruct.update", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "ShapeDtypeStruct.sharding", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "ShapeDtypeStruct.format", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "ShapeDtypeStruct.shape", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapeDtypeStruct.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapeDtypeStruct.weak_type", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapeDtypeStruct.vma", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapeDtypeStruct.is_ref", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapeDtypeStruct.size", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapeDtypeStruct.ndim", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 27, "n_typed": 0, "n_any": 0, "n_untyped": 27, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 8, "n_method_overloads": 8, "n_method_params": 10, "n_attrs": 7, "n_properties": 2 }, { "kind": "class", "name": "jax._src.core.Ref", "methods": [ { "kind": "function", "name": "Ref.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Ref.__repr__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Ref.__getitem__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Ref.__setitem__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Ref.__len__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Ref.addupdate", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Ref.unsafe_buffer_pointer", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "Ref.at", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "Ref._aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Ref._refs", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Ref.aval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Ref.shape", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Ref.size", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Ref.ndim", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Ref.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Ref.sharding", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Ref.format", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Ref.committed", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Ref._committed", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 26, "n_typed": 3, "n_any": 1, "n_untyped": 22, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 7, "n_attrs": 11, "n_properties": 1 }, { "kind": "function", "name": "jax._src.core.free_ref", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.core.ensure_compile_time_eval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.core.typeof", "n_typed": 0, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.core.shaped_abstractify", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax._src.core.Trace", "methods": [ { "kind": "function", "name": "Trace.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Trace.process_primitive", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Trace.invalidate", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Trace.is_valid", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Trace.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Trace.process_call", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "Trace.process_map", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "Trace.process_custom_jvp_call", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "Trace.process_custom_transpose", "n_typed": 2, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "Trace.process_custom_vjp_call", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "Trace.full_raise", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "Trace.main", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "Trace.requires_low", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 41, "n_typed": 2, "n_any": 0, "n_untyped": 39, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 11, "n_method_overloads": 11, "n_method_params": 28, "n_attrs": 1, "n_properties": 1 }, { "kind": "class", "name": "jax._src.core.CallPrimitive", "methods": [ { "kind": "function", "name": "CallPrimitive.bind", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "CallPrimitive.bind_with_trace", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "CallPrimitive.get_bind_params", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "CallPrimitive.multiple_results", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CallPrimitive.call_primitive", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 9, "n_typed": 0, "n_any": 0, "n_untyped": 9, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 6, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax._src.core.get_opaque_trace_state", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.core.eval_context", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "attr", "name": "jax._src.core.no_effects", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.core.no_axis_name", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.core.find_top_trace", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax._src.core.pytype_aval_mappings", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "class", "name": "jax._src.core.Tracer", "methods": [ { "kind": "function", "name": "Tracer.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Tracer._error_repr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.__array__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Tracer._is_traced_array", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.__dlpack__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Tracer.tolist", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.tobytes", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Tracer.full_lower", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.__iter__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.__reversed__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.__len__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.to_concrete_value", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.get_referent", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.__bool__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.__int__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.__float__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.__complex__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.__hex__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.__oct__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.__index__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.__reduce__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.__setitem__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Tracer.__array_module__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Tracer.__getattr__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Tracer._short_repr", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer._pretty_print", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Tracer.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer._contents", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer._origin_msg", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.addressable_data", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Tracer.delete", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.devices", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.is_deleted", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.on_device_size_in_bytes", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Tracer.unsafe_buffer_pointer", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "Tracer.sharding", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Tracer.committed", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Tracer.device", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Tracer.addressable_shards", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Tracer.at", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Tracer.aval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Tracer.block_until_ready", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Tracer.copy_to_host_async", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Tracer.global_shards", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Tracer.is_fully_addressable", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Tracer.is_fully_replicated", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Tracer.traceback", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "Tracer.__array_priority__", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Tracer.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Tracer._trace", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Tracer._line_info", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Tracer.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Tracer.ndim", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Tracer.size", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Tracer.shape", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 65, "n_typed": 7, "n_any": 1, "n_untyped": 57, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 35, "n_method_overloads": 35, "n_method_params": 12, "n_attrs": 8, "n_properties": 12 }, { "kind": "function", "name": "jax._src.core.min_dim", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.core.InconclusiveDimensionOperation", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax._src.core.concretization_function_error", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax._src.core.Atom", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.core.abstract_token", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.core.literalable_types", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax._src.core.new_jaxpr_eqn", "n_typed": 1, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.core.is_constant_dim", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.core.jaxprs_in_params", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax._src.core.ShapedArray", "methods": [ { "kind": "function", "name": "ShapedArray.__init__", "n_typed": 2, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "ShapedArray.lower_val", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ShapedArray.raise_val", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ShapedArray.lo_ty", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ShapedArray.update", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "ShapedArray.__eq__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ShapedArray.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ShapedArray.__ne__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ShapedArray.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ShapedArray.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ShapedArray.to_tangent_aval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ShapedArray.to_cotangent_aval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ShapedArray.str_short", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "ShapedArray._len", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ShapedArray.update_vma", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ShapedArray.update_weak_type", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "ShapedArray.array_abstraction_level", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ShapedArray.shape", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapedArray.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapedArray.weak_type", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapedArray.sharding", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapedArray.vma", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapedArray.memory_space", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapedArray.ndim", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapedArray.size", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapedArray.broadcast", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "ShapedArray.transpose", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "ShapedArray.reshape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "ShapedArray._iter", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "ShapedArray._bool", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapedArray._int", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapedArray._float", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapedArray._complex", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapedArray._hex", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapedArray._oct", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ShapedArray._index", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 54, "n_typed": 6, "n_any": 0, "n_untyped": 48, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 16, "n_method_overloads": 16, "n_method_params": 19, "n_attrs": 20, "n_properties": 0 }, { "kind": "class", "name": "jax._src.core.JaxprTypeError", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.core.JaxprPpSettings", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "JaxprPpSettings.print_shapes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "JaxprPpSettings.source_info", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "JaxprPpSettings.name_stack", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "JaxprPpSettings.custom_pp_eqn_rules", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "JaxprPpSettings.print_effects", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 5, "n_properties": 0 }, { "kind": "function", "name": "jax._src.core.check_jaxpr", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax._src.core.gensym", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "class", "name": "jax._src.core.DropVar", "methods": [ { "kind": "function", "name": "DropVar.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DropVar.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DropVar.pretty_print", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [], "n_typable": 6, "n_typed": 3, "n_any": 0, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 3, "n_attrs": 0, "n_properties": 0 }, { "kind": "attr", "name": "jax._src.core.OutputType", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax._src.core.is_concrete", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax._src.core.aval_mapping_handlers", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax._src.core.primal_dtype_to_tangent_dtype", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.core.valid_jaxtype", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.core.concrete_or_error", "n_typed": 0, "n_any": 2, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax._src.core.trace_ctx", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.core.ParamDict", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax._src.core.eval_jaxpr", "n_typed": 2, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "class", "name": "jax._src.core.AbstractValue", "methods": [ { "kind": "function", "name": "AbstractValue.to_tangent_aval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AbstractValue.to_cotangent_aval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AbstractValue.at_least_vspace", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AbstractValue.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AbstractValue.update_weak_type", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AbstractValue.update_vma", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AbstractValue.strip_weak_type", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AbstractValue.normalize", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AbstractValue.update", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AbstractValue.lo_ty", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AbstractValue.lo_ty_qdd", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AbstractValue.str_short", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "AbstractValue.dec_rank", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "AbstractValue.inc_rank", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "AbstractValue.shard", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "AbstractValue.unshard", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [ { "kind": "property", "name": "AbstractValue.is_high", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "AbstractValue.has_qdd", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [], "n_typable": 35, "n_typed": 4, "n_any": 0, "n_untyped": 31, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 16, "n_method_overloads": 16, "n_method_params": 17, "n_attrs": 0, "n_properties": 2 }, { "kind": "function", "name": "jax._src.core.is_constant_shape", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax._src.core.custom_typechecks", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "class", "name": "jax._src.core.JaxprPpContext", "methods": [ { "kind": "function", "name": "JaxprPpContext.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "JaxprPpContext.suggest_same_var_names", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "JaxprPpContext.var_names", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "JaxprPpContext.shared_jaxprs", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "JaxprPpContext.shared_jaxpr_names", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 7, "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 2, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax._src.core.max_dim", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.core.cur_qdd", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax._src.core.aval_property", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "class", "name": "jax._src.core.AvalMutableQDD", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "AvalMutableQDD.aval", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "AvalMutableQDD.mutable_qdd", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "attr", "name": "jax._src.core.aval_method", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "class", "name": "jax._src.core.AvalQDD", "methods": [ { "kind": "function", "name": "AvalQDD.lo_ty", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AvalQDD.read_loval", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AvalQDD.new_from_loval", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AvalQDD.to_tangent_aval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "AvalQDD.is_high", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "AvalQDD.aval", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "AvalQDD.qdd", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "AvalQDD.has_qdd", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 6, "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 2, "n_attrs": 4, "n_properties": 0 }, { "kind": "class", "name": "jax._src.core.ClosedJaxpr", "methods": [ { "kind": "function", "name": "ClosedJaxpr.__init__", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "ClosedJaxpr.map_jaxpr", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ClosedJaxpr.replace", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "ClosedJaxpr.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ClosedJaxpr.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ClosedJaxpr.pretty_print", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "ClosedJaxpr._repr_pretty_", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [ { "kind": "property", "name": "ClosedJaxpr.in_avals", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "ClosedJaxpr.in_aval_qdds", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "ClosedJaxpr.final_aval_qdds", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "ClosedJaxpr.out_avals", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "ClosedJaxpr._jaxpr", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "ClosedJaxpr._consts", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "ClosedJaxpr.jaxpr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ClosedJaxpr.consts", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ClosedJaxpr.literals", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ClosedJaxpr.constvars", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ClosedJaxpr.invars", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ClosedJaxpr.outvars", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ClosedJaxpr.eqns", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ClosedJaxpr.effects", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ClosedJaxpr.debug_info", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ClosedJaxpr.is_high", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 35, "n_typed": 6, "n_any": 0, "n_untyped": 29, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 13, "n_attrs": 12, "n_properties": 4 }, { "kind": "class", "name": "jax._src.core.Token", "methods": [ { "kind": "function", "name": "Token.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Token.block_until_ready", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Token._buf", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 4, "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 1, "n_attrs": 1, "n_properties": 0 }, { "kind": "class", "name": "jax._src.core.Var", "methods": [ { "kind": "function", "name": "Var.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Var.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Var.pretty_print", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Var.count", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Var.aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Var.initial_qdd", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Var.final_qdd", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 12, "n_typed": 7, "n_any": 0, "n_untyped": 5, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 5, "n_attrs": 4, "n_properties": 0 }, { "kind": "class", "name": "jax._src.core.Jaxpr", "methods": [ { "kind": "function", "name": "Jaxpr.__init__", "n_typed": 7, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "Jaxpr.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Jaxpr.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Jaxpr.pretty_print", "n_typed": 1, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "Jaxpr._repr_pretty_", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Jaxpr.replace", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "Jaxpr.constvars", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Jaxpr.invars", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Jaxpr.outvars", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Jaxpr.eqns", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Jaxpr.effects", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Jaxpr.debug_info", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Jaxpr.is_high", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Jaxpr.in_avals", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Jaxpr.in_aval_qdds", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Jaxpr.final_aval_qdds", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Jaxpr.out_avals", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "Jaxpr._constvars", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Jaxpr._invars", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Jaxpr._outvars", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Jaxpr._eqns", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Jaxpr._effects", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Jaxpr._debug_info", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Jaxpr._is_high", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 40, "n_typed": 24, "n_any": 0, "n_untyped": 16, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 16, "n_attrs": 7, "n_properties": 11 }, { "kind": "function", "name": "jax._src.core.unmapped_aval", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "class", "name": "jax._src.core.JaxprEqn", "methods": [ { "kind": "function", "name": "JaxprEqn.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "JaxprEqn.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "JaxprEqn.replace", "n_typed": 7, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "JaxprEqn.invars", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "JaxprEqn.outvars", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "JaxprEqn.primitive", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "JaxprEqn.params", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "JaxprEqn.effects", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "JaxprEqn.source_info", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "JaxprEqn.ctx", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 24, "n_typed": 14, "n_any": 0, "n_untyped": 10, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 14, "n_attrs": 7, "n_properties": 0 }, { "kind": "function", "name": "jax._src.core.jaxpr_as_fun", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.core.mapped_aval", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax._src.core.Primitive", "methods": [ { "kind": "function", "name": "Primitive.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Primitive.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Primitive.bind", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Primitive._true_bind", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Primitive.bind_with_trace", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Primitive.def_impl", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Primitive.def_abstract_eval", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Primitive.def_effectful_abstract_eval", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Primitive.def_effectful_abstract_eval2", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Primitive.def_bind_with_trace", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Primitive.impl", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Primitive.abstract_eval", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Primitive.get_bind_params", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Primitive.is_high", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Primitive.name", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Primitive.multiple_results", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Primitive.call_primitive", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Primitive.map_primitive", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Primitive.ref_primitive", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Primitive.skip_canonicalization", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Primitive.to_lojax", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Primitive.is_effectful", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 41, "n_typed": 9, "n_any": 0, "n_untyped": 32, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 14, "n_method_overloads": 14, "n_method_params": 20, "n_attrs": 8, "n_properties": 0 }, { "kind": "class", "name": "jax._src.core.Literal", "methods": [ { "kind": "function", "name": "Literal.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Literal.pretty_print", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Literal.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "Literal.hash", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "Literal.val", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Literal.aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Literal.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 10, "n_typed": 3, "n_any": 1, "n_untyped": 6, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 4, "n_attrs": 3, "n_properties": 1 }, { "kind": "attr", "name": "jax._src.core.call_p", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.core.closed_call_p", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "type_ignores": [ { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": [ "missing-attribute" ] }, { "kind": "type", "rules": [ "missing-attribute" ] }, { "kind": "type", "rules": [ "missing-attribute" ] }, { "kind": "type", "rules": [ "missing-attribute" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null } ], "name": "jax._src.core", "names": [ "jax._src.core.AbstractValue", "jax._src.core.Atom", "jax._src.core.AvalMutableQDD", "jax._src.core.AvalQDD", "jax._src.core.CallPrimitive", "jax._src.core.ClosedJaxpr", "jax._src.core.DropVar", "jax._src.core.InconclusiveDimensionOperation", "jax._src.core.Jaxpr", "jax._src.core.JaxprEqn", "jax._src.core.JaxprPpContext", "jax._src.core.JaxprPpSettings", "jax._src.core.JaxprTypeError", "jax._src.core.Literal", "jax._src.core.OutputType", "jax._src.core.ParamDict", "jax._src.core.Primitive", "jax._src.core.Ref", "jax._src.core.ShapeDtypeStruct", "jax._src.core.ShapedArray", "jax._src.core.Token", "jax._src.core.Trace", "jax._src.core.Tracer", "jax._src.core.Var", "jax._src.core.abstract_token", "jax._src.core.aval_mapping_handlers", "jax._src.core.aval_method", "jax._src.core.aval_property", "jax._src.core.call_p", "jax._src.core.check_jaxpr", "jax._src.core.closed_call_p", "jax._src.core.concrete_or_error", "jax._src.core.concretization_function_error", "jax._src.core.cur_qdd", "jax._src.core.custom_typechecks", "jax._src.core.empty_ref", "jax._src.core.ensure_compile_time_eval", "jax._src.core.eval_context", "jax._src.core.eval_jaxpr", "jax._src.core.find_top_trace", "jax._src.core.free_ref", "jax._src.core.freeze", "jax._src.core.gensym", "jax._src.core.get_opaque_trace_state", "jax._src.core.is_concrete", "jax._src.core.is_constant_dim", "jax._src.core.is_constant_shape", "jax._src.core.jaxpr_as_fun", "jax._src.core.jaxprs_in_params", "jax._src.core.literalable_types", "jax._src.core.mapped_aval", "jax._src.core.max_dim", "jax._src.core.min_dim", "jax._src.core.new_jaxpr_eqn", "jax._src.core.no_axis_name", "jax._src.core.no_effects", "jax._src.core.primal_dtype_to_tangent_dtype", "jax._src.core.pytype_aval_mappings", "jax._src.core.shaped_abstractify", "jax._src.core.trace_ctx", "jax._src.core.typeof", "jax._src.core.unmapped_aval", "jax._src.core.valid_jaxtype" ], "n_typable": 524, "n_typed": 132, "n_any": 7, "n_untyped": 385, "n_functions": 26, "n_function_overloads": 26, "n_function_params": 45, "n_methods": 143, "n_method_overloads": 143, "n_method_params": 179, "n_classes": 21, "n_attrs": 121, "n_properties": 34, "n_type_ignores": 18 }, { "path": "jax/__init__.py", "symbol_reports": [ { "kind": "attr", "name": "jax.ds", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [], "name": "jax", "names": [ "jax.ds" ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/ad_checkpoint.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.ad_checkpoint.checkpoint", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "attr", "name": "jax._src.ad_checkpoint.checkpoint_policies", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.ad_checkpoint.remat", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.ad_checkpoint.checkpoint_name", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.ad_checkpoint.print_saved_residuals", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax._src.ad_checkpoint.remat_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.ad_checkpoint.name_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [ { "kind": "type", "rules": null }, { "kind": "type", "rules": null } ], "name": "jax._src.ad_checkpoint", "names": [ "jax._src.ad_checkpoint.checkpoint", "jax._src.ad_checkpoint.checkpoint_name", "jax._src.ad_checkpoint.checkpoint_policies", "jax._src.ad_checkpoint.name_p", "jax._src.ad_checkpoint.print_saved_residuals", "jax._src.ad_checkpoint.remat", "jax._src.ad_checkpoint.remat_p" ], "n_typable": 22, "n_typed": 12, "n_any": 0, "n_untyped": 10, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 15, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 3, "n_properties": 0, "n_type_ignores": 2 }, { "path": "jax/_src/shard_map.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.shard_map.shard_map", "n_typed": 6, "n_any": 1, "n_untyped": 0, "n_overloads": 2, "n_typable": 7, "n_params": 6, "n_function_overloads": 2, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.shard_map.smap", "n_typed": 4, "n_any": 1, "n_untyped": 0, "n_overloads": 2, "n_typable": 5, "n_params": 4, "n_function_overloads": 2, "n_function_params": 4 } ], "type_ignores": [], "name": "jax._src.shard_map", "names": [ "jax._src.shard_map.shard_map", "jax._src.shard_map.smap" ], "n_typable": 12, "n_typed": 10, "n_any": 2, "n_untyped": 0, "n_functions": 2, "n_function_overloads": 4, "n_function_params": 10, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/version.py", "symbol_reports": [ { "kind": "attr", "name": "jax.version.__version_info__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.version.__version__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [], "name": "jax.version", "names": [ "jax.version.__version__", "jax.version.__version_info__" ], "n_typable": 2, "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 2, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/named_sharding.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.named_sharding.NamedSharding", "methods": [ { "kind": "function", "name": "NamedSharding.__init__", "n_typed": 3, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "NamedSharding.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "NamedSharding.__reduce__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "NamedSharding.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "NamedSharding.__eq__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "NamedSharding.check_compatible_aval", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "NamedSharding.with_memory_kind", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "NamedSharding.update", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "NamedSharding._to_xla_hlo_sharding", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "NamedSharding._to_sdy_sharding", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "NamedSharding.memory_kind", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "NamedSharding.num_devices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "NamedSharding.device_set", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "NamedSharding._device_assignment", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "NamedSharding.is_fully_addressable", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "NamedSharding._is_concrete", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "NamedSharding.addressable_devices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "NamedSharding.is_fully_replicated", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "NamedSharding.replicated_axes", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "NamedSharding.mesh", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "NamedSharding.spec", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "NamedSharding._memory_kind", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "NamedSharding._logical_device_ids", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 33, "n_typed": 25, "n_any": 0, "n_untyped": 8, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 10, "n_method_overloads": 10, "n_method_params": 10, "n_attrs": 4, "n_properties": 9 }, { "kind": "class", "name": "jax._src.named_sharding.AUTO", "methods": [ { "kind": "function", "name": "AUTO.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AUTO._to_sdy_sharding", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "AUTO._device_assignment", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "AUTO.mesh", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 6, "n_typed": 3, "n_any": 0, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 2, "n_attrs": 1, "n_properties": 1 } ], "type_ignores": [], "name": "jax._src.named_sharding", "names": [ "jax._src.named_sharding.AUTO", "jax._src.named_sharding.NamedSharding" ], "n_typable": 39, "n_typed": 28, "n_any": 0, "n_untyped": 11, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 12, "n_method_overloads": 12, "n_method_params": 12, "n_classes": 2, "n_attrs": 5, "n_properties": 10, "n_type_ignores": 0 }, { "path": "jax/_src/ref.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.ref.new_ref", "n_typed": 1, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [], "name": "jax._src.ref", "names": [ "jax._src.ref.new_ref" ], "n_typable": 3, "n_typed": 1, "n_any": 2, "n_untyped": 0, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 2, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/environment_info.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.environment_info.print_environment_info", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax._src.environment_info", "names": [ "jax._src.environment_info.print_environment_info" ], "n_typable": 2, "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 1, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/interpreters/partial_eval.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.interpreters.partial_eval.Offloadable", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "Offloadable.src", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Offloadable.dst", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "attr", "name": "jax._src.interpreters.partial_eval.Saveable", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.interpreters.partial_eval.Recompute", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.interpreters.partial_eval.trace_to_jaxpr_nounits", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.interpreters.partial_eval.dce_jaxpr_closed_call_rule", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.interpreters.partial_eval.dce_jaxpr_consts", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.interpreters.partial_eval.dce_jaxpr_call_rule", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.interpreters.partial_eval.PartialVal", "methods": [ { "kind": "function", "name": "PartialVal.__new__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "PartialVal.known", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "PartialVal.unknown", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "PartialVal.is_known", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "PartialVal.get_known", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "PartialVal.get_aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [], "n_typable": 9, "n_typed": 8, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 3, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax._src.interpreters.partial_eval.trace_to_jaxpr_dynamic", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "attr", "name": "jax._src.interpreters.partial_eval.partial_eval_jaxpr_custom_rules", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "class", "name": "jax._src.interpreters.partial_eval.DynamicJaxprTracer", "methods": [ { "kind": "function", "name": "DynamicJaxprTracer.__init__", "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "DynamicJaxprTracer._short_repr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DynamicJaxprTracer.cur_qdd", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DynamicJaxprTracer.full_lower", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DynamicJaxprTracer._contents", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DynamicJaxprTracer._origin_msg", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DynamicJaxprTracer.get_const", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DynamicJaxprTracer.get_referent", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "DynamicJaxprTracer.aval_mutable_qdd", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "DynamicJaxprTracer._trace", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "DynamicJaxprTracer.aval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "DynamicJaxprTracer.val", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "DynamicJaxprTracer.mutable_qdd", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "DynamicJaxprTracer.parent", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 19, "n_typed": 6, "n_any": 0, "n_untyped": 13, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 8, "n_method_overloads": 8, "n_method_params": 5, "n_attrs": 5, "n_properties": 1 }, { "kind": "attr", "name": "jax._src.interpreters.partial_eval.custom_partial_eval_rules", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "class", "name": "jax._src.interpreters.partial_eval.JaxprTracer", "methods": [ { "kind": "function", "name": "JaxprTracer.__init__", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "JaxprTracer.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "JaxprTracer.full_lower", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "JaxprTracer.is_known", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "JaxprTracer.get_referent", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "JaxprTracer.aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "JaxprTracer.parents", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "JaxprTracer._trace", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "JaxprTracer.pval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "JaxprTracer.recipe", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 13, "n_typed": 6, "n_any": 0, "n_untyped": 7, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 3, "n_attrs": 3, "n_properties": 2 }, { "kind": "attr", "name": "jax._src.interpreters.partial_eval.dce_rules", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax._src.interpreters.partial_eval.dce_jaxpr", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [ { "kind": "type", "rules": null }, { "kind": "pyrefly", "rules": [ "bad-param-name-override" ] }, { "kind": "pyrefly", "rules": [ "bad-param-name-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-param-name-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] } ], "name": "jax._src.interpreters.partial_eval", "names": [ "jax._src.interpreters.partial_eval.DynamicJaxprTracer", "jax._src.interpreters.partial_eval.JaxprTracer", "jax._src.interpreters.partial_eval.Offloadable", "jax._src.interpreters.partial_eval.PartialVal", "jax._src.interpreters.partial_eval.Recompute", "jax._src.interpreters.partial_eval.Saveable", "jax._src.interpreters.partial_eval.custom_partial_eval_rules", "jax._src.interpreters.partial_eval.dce_jaxpr", "jax._src.interpreters.partial_eval.dce_jaxpr_call_rule", "jax._src.interpreters.partial_eval.dce_jaxpr_closed_call_rule", "jax._src.interpreters.partial_eval.dce_jaxpr_consts", "jax._src.interpreters.partial_eval.dce_rules", "jax._src.interpreters.partial_eval.partial_eval_jaxpr_custom_rules", "jax._src.interpreters.partial_eval.trace_to_jaxpr_dynamic", "jax._src.interpreters.partial_eval.trace_to_jaxpr_nounits" ], "n_typable": 70, "n_typed": 47, "n_any": 0, "n_untyped": 23, "n_functions": 6, "n_function_overloads": 6, "n_function_params": 18, "n_methods": 19, "n_method_overloads": 19, "n_method_params": 11, "n_classes": 4, "n_attrs": 15, "n_properties": 3, "n_type_ignores": 7 }, { "path": "jax/_src/api_util.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.api_util.flatten_axes", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.api_util.argnums_partial", "n_typed": 3, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.api_util.flatten_fun", "n_typed": 3, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.api_util.donation_vector", "n_typed": 2, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.api_util.debug_info", "n_typed": 10, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "jax._src.api_util.rebase_donate_argnums", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.api_util.flatten_fun_nokwargs", "n_typed": 3, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax._src.api_util", "names": [ "jax._src.api_util.argnums_partial", "jax._src.api_util.debug_info", "jax._src.api_util.donation_vector", "jax._src.api_util.flatten_axes", "jax._src.api_util.flatten_fun", "jax._src.api_util.flatten_fun_nokwargs", "jax._src.api_util.rebase_donate_argnums" ], "n_typable": 39, "n_typed": 22, "n_any": 0, "n_untyped": 17, "n_functions": 7, "n_function_overloads": 7, "n_function_params": 32, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/util.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.util.safe_map", "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_overloads": 4, "n_typable": 7, "n_params": 6, "n_function_overloads": 4, "n_function_params": 6 } ], "type_ignores": [ { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] } ], "name": "jax._src.util", "names": [ "jax._src.util.safe_map" ], "n_typable": 7, "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_functions": 1, "n_function_overloads": 4, "n_function_params": 6, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 2 }, { "path": "jax/collect_profile.py", "symbol_reports": [ { "kind": "attr", "name": "jax.collect_profile.known_args", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.collect_profile.main", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.collect_profile.parser", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.collect_profile.collect_profile", "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "attr", "name": "jax.collect_profile.DEFAULT_NUM_TRACING_ATTEMPTS", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.collect_profile.unknown_flags", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [], "name": "jax.collect_profile", "names": [ "jax.collect_profile.DEFAULT_NUM_TRACING_ATTEMPTS", "jax.collect_profile.collect_profile", "jax.collect_profile.known_args", "jax.collect_profile.main", "jax.collect_profile.parser", "jax.collect_profile.unknown_flags" ], "n_typable": 13, "n_typed": 6, "n_any": 0, "n_untyped": 7, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 8, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 4, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/effects.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.effects.Effect", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "attr", "name": "jax._src.effects.Effects", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.effects.control_flow_allowed_effects", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.effects.lowerable_effects", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "type_ignores": [], "name": "jax._src.effects", "names": [ "jax._src.effects.Effect", "jax._src.effects.Effects", "jax._src.effects.control_flow_allowed_effects", "jax._src.effects.lowerable_effects" ], "n_typable": 3, "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 3, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/core.py", "symbol_reports": [ { "kind": "attr", "name": "jax.core.call", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.core.Value", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [], "name": "jax.core", "names": [ "jax.core.Value", "jax.core.call" ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 2, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/linear_util.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.linear_util.DebugInfo", "methods": [ { "kind": "function", "name": "DebugInfo.resolve_result_paths", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DebugInfo.replace_func_name", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DebugInfo.set_result_paths", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DebugInfo.safe_arg_names", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DebugInfo.assert_arg_names", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DebugInfo.filter_arg_names", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DebugInfo.safe_result_paths", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DebugInfo.assert_result_paths", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DebugInfo.filter_result_paths", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DebugInfo.with_unknown_names", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "DebugInfo.func_name", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "DebugInfo.func_filename", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "DebugInfo.func_lineno", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "DebugInfo.traced_for", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DebugInfo.func_src_info", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DebugInfo.arg_names", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DebugInfo.result_paths", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 21, "n_typed": 17, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 10, "n_method_overloads": 10, "n_method_params": 8, "n_attrs": 4, "n_properties": 3 }, { "kind": "function", "name": "jax._src.linear_util.transformation_with_aux", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.linear_util.transformation_with_aux2", "n_typed": 4, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.linear_util.transformation", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax._src.linear_util.StoreException", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax._src.linear_util.cache", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.linear_util.transformation2", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.linear_util.merge_linear_aux", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.linear_util.WrappedFun", "methods": [ { "kind": "function", "name": "WrappedFun.__init__", "n_typed": 7, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "WrappedFun.wrap", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "WrappedFun.populate_stores", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "WrappedFun.call_wrapped", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "WrappedFun.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "WrappedFun.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "WrappedFun.__eq__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "WrappedFun.replace_debug_info", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "WrappedFun.with_unknown_names", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "WrappedFun.__name__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "WrappedFun.f", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "WrappedFun.f_transformed", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "WrappedFun.transforms", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "WrappedFun.stores", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "WrappedFun.params", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "WrappedFun.in_type", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "WrappedFun.debug_info", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 32, "n_typed": 19, "n_any": 0, "n_untyped": 13, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 9, "n_method_overloads": 9, "n_method_params": 15, "n_attrs": 7, "n_properties": 1 } ], "type_ignores": [], "name": "jax._src.linear_util", "names": [ "jax._src.linear_util.DebugInfo", "jax._src.linear_util.StoreException", "jax._src.linear_util.WrappedFun", "jax._src.linear_util.cache", "jax._src.linear_util.merge_linear_aux", "jax._src.linear_util.transformation", "jax._src.linear_util.transformation2", "jax._src.linear_util.transformation_with_aux", "jax._src.linear_util.transformation_with_aux2" ], "n_typable": 77, "n_typed": 48, "n_any": 0, "n_untyped": 29, "n_functions": 6, "n_function_overloads": 6, "n_function_params": 18, "n_methods": 19, "n_method_overloads": 19, "n_method_params": 23, "n_classes": 3, "n_attrs": 11, "n_properties": 4, "n_type_ignores": 0 }, { "path": "jax/_src/custom_batching.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.custom_batching.sequential_vmap", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax._src.custom_batching.custom_vmap", "methods": [ { "kind": "function", "name": "custom_vmap.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "custom_vmap.def_vmap", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "custom_vmap.__call__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "custom_vmap.fun", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "custom_vmap.vmap_rule", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "custom_vmap.__getattr__", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 9, "n_typed": 5, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 4, "n_attrs": 3, "n_properties": 0 } ], "type_ignores": [], "name": "jax._src.custom_batching", "names": [ "jax._src.custom_batching.custom_vmap", "jax._src.custom_batching.sequential_vmap" ], "n_typable": 11, "n_typed": 5, "n_any": 0, "n_untyped": 6, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 1, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 4, "n_classes": 1, "n_attrs": 3, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/ad_util.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.ad_util.zero_from_primal", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.ad_util.SymbolicZero", "methods": [ { "kind": "function", "name": "SymbolicZero.__init__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SymbolicZero.__repr__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "SymbolicZero.__getattr__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SymbolicZero.from_primal_value", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "SymbolicZero.aval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 8, "n_typed": 4, "n_any": 1, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 3, "n_attrs": 1, "n_properties": 0 }, { "kind": "class", "name": "jax._src.ad_util.Zero", "methods": [ { "kind": "function", "name": "Zero.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Zero.__repr__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Zero.instantiate", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Zero.aval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 5, "n_typed": 2, "n_any": 0, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 1, "n_attrs": 1, "n_properties": 0 }, { "kind": "attr", "name": "jax._src.ad_util.stop_gradient_p", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.ad_util.add_jaxvals_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.ad_util.add_jaxvals", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.ad_util.zeros_like_aval", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [ { "kind": "type", "rules": null } ], "name": "jax._src.ad_util", "names": [ "jax._src.ad_util.SymbolicZero", "jax._src.ad_util.Zero", "jax._src.ad_util.add_jaxvals", "jax._src.ad_util.add_jaxvals_p", "jax._src.ad_util.stop_gradient_p", "jax._src.ad_util.zero_from_primal", "jax._src.ad_util.zeros_like_aval" ], "n_typable": 23, "n_typed": 12, "n_any": 1, "n_untyped": 10, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 5, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 4, "n_classes": 2, "n_attrs": 4, "n_properties": 0, "n_type_ignores": 1 }, { "path": "jax/_src/custom_transpose.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.custom_transpose.custom_transpose", "methods": [ { "kind": "function", "name": "custom_transpose.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "custom_transpose.def_transpose", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "custom_transpose.__call__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "custom_transpose.fun", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "custom_transpose.transpose", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "custom_transpose.__getattr__", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 10, "n_typed": 4, "n_any": 0, "n_untyped": 6, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 5, "n_attrs": 3, "n_properties": 0 } ], "type_ignores": [], "name": "jax._src.custom_transpose", "names": [ "jax._src.custom_transpose.custom_transpose" ], "n_typable": 10, "n_typed": 4, "n_any": 0, "n_untyped": 6, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 5, "n_classes": 1, "n_attrs": 3, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/debugging.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.debugging.debug_print", "n_typed": 6, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "attr", "name": "jax._src.debugging.debug_log", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.debugging.visualize_sharding", "n_typed": 7, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.debugging.visualize_array_sharding", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.debugging.DebugEffect", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "DebugEffect.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 1, "n_properties": 0 }, { "kind": "function", "name": "jax._src.debugging.debug_callback", "n_typed": 4, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.debugging.inspect_array_sharding", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [ { "kind": "type", "rules": null } ], "name": "jax._src.debugging", "names": [ "jax._src.debugging.DebugEffect", "jax._src.debugging.debug_callback", "jax._src.debugging.debug_log", "jax._src.debugging.debug_print", "jax._src.debugging.inspect_array_sharding", "jax._src.debugging.visualize_array_sharding", "jax._src.debugging.visualize_sharding" ], "n_typable": 29, "n_typed": 18, "n_any": 2, "n_untyped": 9, "n_functions": 5, "n_function_overloads": 5, "n_function_params": 23, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 2, "n_properties": 0, "n_type_ignores": 1 }, { "path": "jax/_src/debugger/core.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.debugger.core.breakpoint", "n_typed": 4, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 } ], "type_ignores": [], "name": "jax._src.debugger.core", "names": [ "jax._src.debugger.core.breakpoint" ], "n_typable": 7, "n_typed": 4, "n_any": 0, "n_untyped": 3, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 6, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/distributed.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.distributed.is_initialized", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.distributed.initialize", "n_typed": 11, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 12, "n_params": 11, "n_function_overloads": 1, "n_function_params": 11 }, { "kind": "function", "name": "jax._src.distributed.shutdown", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "type_ignores": [], "name": "jax._src.distributed", "names": [ "jax._src.distributed.initialize", "jax._src.distributed.is_initialized", "jax._src.distributed.shutdown" ], "n_typable": 14, "n_typed": 12, "n_any": 0, "n_untyped": 2, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 11, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/dlpack.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.dlpack.from_dlpack", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.dlpack.is_supported_dtype", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax._src.dlpack", "names": [ "jax._src.dlpack.from_dlpack", "jax._src.dlpack.is_supported_dtype" ], "n_typable": 6, "n_typed": 4, "n_any": 0, "n_untyped": 2, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 4, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/errors.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.errors.TracerIntegerConversionError", "methods": [ { "kind": "function", "name": "TracerIntegerConversionError.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [], "n_typable": 2, "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 1, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.errors.JAXIndexError", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.errors.UnexpectedTracerError", "methods": [ { "kind": "function", "name": "UnexpectedTracerError.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [], "n_typable": 2, "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 1, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.errors.TracerBoolConversionError", "methods": [ { "kind": "function", "name": "TracerBoolConversionError.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [], "n_typable": 2, "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 1, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.errors.KeyReuseError", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.errors.NonConcreteBooleanIndexError", "methods": [ { "kind": "function", "name": "NonConcreteBooleanIndexError.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [], "n_typable": 2, "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 1, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.errors.JAXTypeError", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.errors.ConcretizationTypeError", "methods": [ { "kind": "function", "name": "ConcretizationTypeError.__init__", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [], "n_typable": 3, "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 2, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.errors.TracerArrayConversionError", "methods": [ { "kind": "function", "name": "TracerArrayConversionError.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [], "n_typable": 2, "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 1, "n_attrs": 0, "n_properties": 0 } ], "type_ignores": [], "name": "jax._src.errors", "names": [ "jax._src.errors.ConcretizationTypeError", "jax._src.errors.JAXIndexError", "jax._src.errors.JAXTypeError", "jax._src.errors.KeyReuseError", "jax._src.errors.NonConcreteBooleanIndexError", "jax._src.errors.TracerArrayConversionError", "jax._src.errors.TracerBoolConversionError", "jax._src.errors.TracerIntegerConversionError", "jax._src.errors.UnexpectedTracerError" ], "n_typable": 13, "n_typed": 7, "n_any": 0, "n_untyped": 6, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 7, "n_classes": 9, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/example_libraries/optimizers.py", "symbol_reports": [ { "kind": "function", "name": "jax.example_libraries.optimizers.rmsprop", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.example_libraries.optimizers.sm3", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.example_libraries.optimizers.inverse_time_decay", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.example_libraries.optimizers.optimizer", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax.example_libraries.optimizers.JoinPoint", "methods": [ { "kind": "function", "name": "JoinPoint.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "JoinPoint.__iter__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "JoinPoint.subtree", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 4, "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 1, "n_attrs": 1, "n_properties": 0 }, { "kind": "function", "name": "jax.example_libraries.optimizers.unpack_optimizer_state", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.example_libraries.optimizers.ParamsFn", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.example_libraries.optimizers.Schedule", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.example_libraries.optimizers.l2_norm", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.example_libraries.optimizers.OptimizerState", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.example_libraries.optimizers.Step", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.example_libraries.optimizers.UpdateFn", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.example_libraries.optimizers.pack_optimizer_state", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.example_libraries.optimizers.momentum", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.example_libraries.optimizers.make_schedule", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.example_libraries.optimizers.polynomial_decay", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.example_libraries.optimizers.constant", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.example_libraries.optimizers.adagrad", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.example_libraries.optimizers.adam", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.example_libraries.optimizers.rmsprop_momentum", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.example_libraries.optimizers.exponential_decay", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.example_libraries.optimizers.nesterov", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax.example_libraries.optimizers.Optimizer", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "Optimizer.init_fn", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Optimizer.update_fn", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Optimizer.params_fn", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax.example_libraries.optimizers.sgd", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.example_libraries.optimizers.clip_grads", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.example_libraries.optimizers.adamax", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "attr", "name": "jax.example_libraries.optimizers.InitFn", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.example_libraries.optimizers.piecewise_constant", "n_typed": 0, "n_any": 2, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [], "name": "jax.example_libraries.optimizers", "names": [ "jax.example_libraries.optimizers.InitFn", "jax.example_libraries.optimizers.JoinPoint", "jax.example_libraries.optimizers.Optimizer", "jax.example_libraries.optimizers.OptimizerState", "jax.example_libraries.optimizers.ParamsFn", "jax.example_libraries.optimizers.Schedule", "jax.example_libraries.optimizers.Step", "jax.example_libraries.optimizers.UpdateFn", "jax.example_libraries.optimizers.adagrad", "jax.example_libraries.optimizers.adam", "jax.example_libraries.optimizers.adamax", "jax.example_libraries.optimizers.clip_grads", "jax.example_libraries.optimizers.constant", "jax.example_libraries.optimizers.exponential_decay", "jax.example_libraries.optimizers.inverse_time_decay", "jax.example_libraries.optimizers.l2_norm", "jax.example_libraries.optimizers.make_schedule", "jax.example_libraries.optimizers.momentum", "jax.example_libraries.optimizers.nesterov", "jax.example_libraries.optimizers.optimizer", "jax.example_libraries.optimizers.pack_optimizer_state", "jax.example_libraries.optimizers.piecewise_constant", "jax.example_libraries.optimizers.polynomial_decay", "jax.example_libraries.optimizers.rmsprop", "jax.example_libraries.optimizers.rmsprop_momentum", "jax.example_libraries.optimizers.sgd", "jax.example_libraries.optimizers.sm3", "jax.example_libraries.optimizers.unpack_optimizer_state" ], "n_typable": 73, "n_typed": 13, "n_any": 2, "n_untyped": 58, "n_functions": 20, "n_function_overloads": 20, "n_function_params": 45, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 1, "n_classes": 2, "n_attrs": 10, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/example_libraries/stax.py", "symbol_reports": [ { "kind": "attr", "name": "jax.example_libraries.stax.SumPool", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.example_libraries.stax.FanInSum", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax.example_libraries.stax.GeneralConv", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "attr", "name": "jax.example_libraries.stax.Sigmoid", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.example_libraries.stax.BatchNorm", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.example_libraries.stax.FanOut", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.example_libraries.stax.Flatten", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "attr", "name": "jax.example_libraries.stax.Relu", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.example_libraries.stax.Conv", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.example_libraries.stax.Exp", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.example_libraries.stax.parallel", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.example_libraries.stax.ConvTranspose", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.example_libraries.stax.Dropout", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.example_libraries.stax.Softmax", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.example_libraries.stax.Elu", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.example_libraries.stax.Identity", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "attr", "name": "jax.example_libraries.stax.Conv1DTranspose", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.example_libraries.stax.Gelu", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.example_libraries.stax.Tanh", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.example_libraries.stax.shape_dependent", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.example_libraries.stax.Selu", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.example_libraries.stax.elementwise", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.example_libraries.stax.serial", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.example_libraries.stax.MaxPool", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.example_libraries.stax.FanInConcat", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.example_libraries.stax.Softplus", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.example_libraries.stax.AvgPool", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.example_libraries.stax.LogSoftmax", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.example_libraries.stax.LeakyRelu", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.example_libraries.stax.Dense", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.example_libraries.stax.GeneralConvTranspose", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 } ], "type_ignores": [], "name": "jax.example_libraries.stax", "names": [ "jax.example_libraries.stax.AvgPool", "jax.example_libraries.stax.BatchNorm", "jax.example_libraries.stax.Conv", "jax.example_libraries.stax.Conv1DTranspose", "jax.example_libraries.stax.ConvTranspose", "jax.example_libraries.stax.Dense", "jax.example_libraries.stax.Dropout", "jax.example_libraries.stax.Elu", "jax.example_libraries.stax.Exp", "jax.example_libraries.stax.FanInConcat", "jax.example_libraries.stax.FanInSum", "jax.example_libraries.stax.FanOut", "jax.example_libraries.stax.Flatten", "jax.example_libraries.stax.Gelu", "jax.example_libraries.stax.GeneralConv", "jax.example_libraries.stax.GeneralConvTranspose", "jax.example_libraries.stax.Identity", "jax.example_libraries.stax.LeakyRelu", "jax.example_libraries.stax.LogSoftmax", "jax.example_libraries.stax.MaxPool", "jax.example_libraries.stax.Relu", "jax.example_libraries.stax.Selu", "jax.example_libraries.stax.Sigmoid", "jax.example_libraries.stax.Softmax", "jax.example_libraries.stax.Softplus", "jax.example_libraries.stax.SumPool", "jax.example_libraries.stax.Tanh", "jax.example_libraries.stax.elementwise", "jax.example_libraries.stax.parallel", "jax.example_libraries.stax.serial", "jax.example_libraries.stax.shape_dependent" ], "n_typable": 63, "n_typed": 0, "n_any": 0, "n_untyped": 63, "n_functions": 14, "n_function_overloads": 14, "n_function_params": 32, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 17, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/earray.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.earray.EArray", "methods": [ { "kind": "function", "name": "EArray.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "EArray.block_until_ready", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "EArray.copy_to_host_async", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "EArray.copy", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "EArray.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "EArray.__iter__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "EArray.__len__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "EArray.addressable_data", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "EArray.sharding", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "EArray.committed", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "EArray.device", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "EArray.addressable_shards", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "EArray.global_shards", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "EArray.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "EArray.__array_priority__", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "EArray.aval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "EArray.shape", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "EArray.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "EArray.ndim", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "EArray.size", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "EArray.itemsize", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "EArray.devices", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "EArray._committed", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "EArray.is_fully_addressable", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "EArray.is_fully_replicated", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "EArray.delete", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "EArray.is_deleted", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "EArray.on_device_size_in_bytes", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "EArray.unsafe_buffer_pointer", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 30, "n_typed": 2, "n_any": 0, "n_untyped": 28, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 8, "n_method_overloads": 8, "n_method_params": 3, "n_attrs": 16, "n_properties": 5 } ], "type_ignores": [ { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] } ], "name": "jax._src.earray", "names": [ "jax._src.earray.EArray" ], "n_typable": 30, "n_typed": 2, "n_any": 0, "n_untyped": 28, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 8, "n_method_overloads": 8, "n_method_params": 3, "n_classes": 1, "n_attrs": 16, "n_properties": 5, "n_type_ignores": 13 }, { "path": "jax/experimental/array_serialization/pytree_serialization.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.array_serialization.pytree_serialization.nonblocking_save", "n_typed": 4, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.array_serialization.pytree_serialization.load_pytreedef", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.array_serialization.pytree_serialization.load", "n_typed": 3, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.array_serialization.pytree_serialization.save", "n_typed": 4, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.array_serialization.pytree_serialization.nonblocking_load", "n_typed": 4, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax.experimental.array_serialization.pytree_serialization", "names": [ "jax.experimental.array_serialization.pytree_serialization.load", "jax.experimental.array_serialization.pytree_serialization.load_pytreedef", "jax.experimental.array_serialization.pytree_serialization.nonblocking_load", "jax.experimental.array_serialization.pytree_serialization.nonblocking_save", "jax.experimental.array_serialization.pytree_serialization.save" ], "n_typable": 22, "n_typed": 16, "n_any": 6, "n_untyped": 0, "n_functions": 5, "n_function_overloads": 5, "n_function_params": 17, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/array_serialization/pytree_serialization_utils.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.array_serialization.pytree_serialization_utils.serialize_pytreedef", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.array_serialization.pytree_serialization_utils.deserialize_pytreedef", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.array_serialization.pytree_serialization_utils", "names": [ "jax.experimental.array_serialization.pytree_serialization_utils.deserialize_pytreedef", "jax.experimental.array_serialization.pytree_serialization_utils.serialize_pytreedef" ], "n_typable": 4, "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 2, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/export/_export.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.export._export.register_pytree_node_serialization", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "class", "name": "jax._src.export._export.DisabledSafetyCheck", "methods": [ { "kind": "function", "name": "DisabledSafetyCheck.platform", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DisabledSafetyCheck.custom_call", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DisabledSafetyCheck.is_custom_call", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DisabledSafetyCheck.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DisabledSafetyCheck.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DisabledSafetyCheck.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DisabledSafetyCheck.__eq__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DisabledSafetyCheck.__hash__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "DisabledSafetyCheck._impl", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 12, "n_typed": 8, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 8, "n_method_overloads": 8, "n_method_params": 3, "n_attrs": 1, "n_properties": 0 }, { "kind": "attr", "name": "jax._src.export._export.minimum_supported_calling_convention_version", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax._src.export._export.register_namedtuple_serialization", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.export._export.default_export_platform", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.export._export.deserialize", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax._src.export._export.Exported", "methods": [ { "kind": "function", "name": "Exported.mlir_module", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Exported.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Exported.in_shardings_jax", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Exported.out_shardings_jax", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Exported.has_vjp", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Exported.vjp", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Exported.serialize", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Exported.call", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Exported.fun_name", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported.in_tree", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported.in_avals", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported.out_tree", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported.out_avals", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported._has_named_shardings", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported._in_named_shardings", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported._out_named_shardings", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported.in_shardings_hlo", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported.out_shardings_hlo", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported.nr_devices", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported.platforms", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported.ordered_effects", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported.unordered_effects", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported.disabled_safety_checks", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported.mlir_module_serialized", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported.calling_convention_version", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported.module_kept_var_idx", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported.uses_global_constants", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Exported._get_vjp", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 13, "n_typed": 9, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 8, "n_method_overloads": 8, "n_method_params": 5, "n_attrs": 20, "n_properties": 0 }, { "kind": "function", "name": "jax._src.export._export.export", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "attr", "name": "jax._src.export._export.maximum_supported_calling_convention_version", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [], "name": "jax._src.export._export", "names": [ "jax._src.export._export.DisabledSafetyCheck", "jax._src.export._export.Exported", "jax._src.export._export.default_export_platform", "jax._src.export._export.deserialize", "jax._src.export._export.export", "jax._src.export._export.maximum_supported_calling_convention_version", "jax._src.export._export.minimum_supported_calling_convention_version", "jax._src.export._export.register_namedtuple_serialization", "jax._src.export._export.register_pytree_node_serialization" ], "n_typable": 42, "n_typed": 34, "n_any": 0, "n_untyped": 8, "n_functions": 5, "n_function_overloads": 5, "n_function_params": 12, "n_methods": 16, "n_method_overloads": 16, "n_method_params": 8, "n_classes": 2, "n_attrs": 23, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/array_serialization/serialization.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.array_serialization.serialization.logger", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "class", "name": "jax.experimental.array_serialization.serialization.GlobalAsyncCheckpointManager", "methods": [ { "kind": "function", "name": "GlobalAsyncCheckpointManager.serialize", "n_typed": 2, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "GlobalAsyncCheckpointManager.serialize_with_paths", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "GlobalAsyncCheckpointManager.deserialize", "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "GlobalAsyncCheckpointManager.deserialize_with_paths", "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "properties": [], "attrs": [], "n_typable": 22, "n_typed": 16, "n_any": 0, "n_untyped": 6, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 18, "n_attrs": 0, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.array_serialization.serialization.get_tensorstore_spec", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "class", "name": "jax.experimental.array_serialization.serialization.GlobalAsyncCheckpointManagerBase", "methods": [ { "kind": "function", "name": "GlobalAsyncCheckpointManagerBase.check_for_errors", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "GlobalAsyncCheckpointManagerBase.wait_until_finished", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "GlobalAsyncCheckpointManagerBase.serialize", "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "GlobalAsyncCheckpointManagerBase.deserialize", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "properties": [], "attrs": [], "n_typable": 11, "n_typed": 5, "n_any": 0, "n_untyped": 6, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 7, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.array_serialization.serialization.is_remote_storage", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax.experimental.array_serialization.serialization.BarrierTimeoutError", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.array_serialization.serialization.AsyncManager", "methods": [ { "kind": "function", "name": "AsyncManager.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AsyncManager.__del__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AsyncManager._thread_func", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AsyncManager._start_async_commit", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AsyncManager.check_for_errors", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AsyncManager.wait_until_finished", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AsyncManager._add_futures", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [], "n_typable": 10, "n_typed": 1, "n_any": 0, "n_untyped": 9, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 3, "n_attrs": 0, "n_properties": 0 } ], "type_ignores": [], "name": "jax.experimental.array_serialization.serialization", "names": [ "jax.experimental.array_serialization.serialization.AsyncManager", "jax.experimental.array_serialization.serialization.BarrierTimeoutError", "jax.experimental.array_serialization.serialization.GlobalAsyncCheckpointManager", "jax.experimental.array_serialization.serialization.GlobalAsyncCheckpointManagerBase", "jax.experimental.array_serialization.serialization.get_tensorstore_spec", "jax.experimental.array_serialization.serialization.is_remote_storage", "jax.experimental.array_serialization.serialization.logger" ], "n_typable": 47, "n_typed": 24, "n_any": 0, "n_untyped": 23, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 1, "n_methods": 15, "n_method_overloads": 15, "n_method_params": 28, "n_classes": 4, "n_attrs": 2, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/array_serialization/tensorstore_impl.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.array_serialization.tensorstore_impl.verify_tensorstore_spec", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.experimental.array_serialization.tensorstore_impl.async_deserialize", "n_typed": 5, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.experimental.array_serialization.tensorstore_impl.get_tensorstore_spec", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "attr", "name": "jax.experimental.array_serialization.tensorstore_impl.logger", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.array_serialization.tensorstore_impl.merge_nested_ts_specs", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.array_serialization.tensorstore_impl.estimate_read_memory_footprint", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.array_serialization.tensorstore_impl.async_serialize", "n_typed": 3, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "attr", "name": "jax.experimental.array_serialization.tensorstore_impl.Future", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.array_serialization.tensorstore_impl.Transaction", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.array_serialization.tensorstore_impl.is_tensorstore_spec_leaf", "n_typed": 0, "n_any": 1, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.array_serialization.tensorstore_impl.combine_kvstores", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax.experimental.array_serialization.tensorstore_impl", "names": [ "jax.experimental.array_serialization.tensorstore_impl.Future", "jax.experimental.array_serialization.tensorstore_impl.Transaction", "jax.experimental.array_serialization.tensorstore_impl.async_deserialize", "jax.experimental.array_serialization.tensorstore_impl.async_serialize", "jax.experimental.array_serialization.tensorstore_impl.combine_kvstores", "jax.experimental.array_serialization.tensorstore_impl.estimate_read_memory_footprint", "jax.experimental.array_serialization.tensorstore_impl.get_tensorstore_spec", "jax.experimental.array_serialization.tensorstore_impl.is_tensorstore_spec_leaf", "jax.experimental.array_serialization.tensorstore_impl.logger", "jax.experimental.array_serialization.tensorstore_impl.merge_nested_ts_specs", "jax.experimental.array_serialization.tensorstore_impl.verify_tensorstore_spec" ], "n_typable": 43, "n_typed": 29, "n_any": 1, "n_untyped": 13, "n_functions": 8, "n_function_overloads": 8, "n_function_params": 34, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 3, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/buffer_callback.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.buffer_callback.ExecutionContext", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.buffer_callback.ExecutionStage", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.buffer_callback.buffer_callback", "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "attr", "name": "jax._src.buffer_callback.Buffer", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [], "name": "jax._src.buffer_callback", "names": [ "jax._src.buffer_callback.Buffer", "jax._src.buffer_callback.ExecutionContext", "jax._src.buffer_callback.ExecutionStage", "jax._src.buffer_callback.buffer_callback" ], "n_typable": 10, "n_typed": 6, "n_any": 0, "n_untyped": 4, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 6, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 3, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/checkify.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.checkify.Error", "methods": [ { "kind": "function", "name": "Error.get", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Error.get_exception", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Error.throw", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Error.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Error._get_batched_exception", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Error._update", "n_typed": 1, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "Error._add_placeholder_effects", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Error._replace", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Error.tree_flatten", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Error.tree_unflatten", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Error._pred", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Error._code", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Error._metadata", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Error._payload", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 20, "n_typed": 5, "n_any": 0, "n_untyped": 15, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 10, "n_method_overloads": 10, "n_method_params": 10, "n_attrs": 4, "n_properties": 0 }, { "kind": "attr", "name": "jax._src.checkify.user_checks", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.checkify.checkify", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.checkify.check", "n_typed": 4, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "class", "name": "jax._src.checkify.JaxRuntimeError", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "attr", "name": "jax._src.checkify.ErrorCategory", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax._src.checkify.float_checks", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax._src.checkify.all_checks", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax._src.checkify.automatic_checks", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax._src.checkify.init_error", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.checkify.debug_check", "n_typed": 3, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "attr", "name": "jax._src.checkify.index_checks", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.checkify.nan_checks", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.checkify.check_error", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax._src.checkify.div_checks", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [ { "kind": "type", "rules": null } ], "name": "jax._src.checkify", "names": [ "jax._src.checkify.Error", "jax._src.checkify.ErrorCategory", "jax._src.checkify.JaxRuntimeError", "jax._src.checkify.all_checks", "jax._src.checkify.automatic_checks", "jax._src.checkify.check", "jax._src.checkify.check_error", "jax._src.checkify.checkify", "jax._src.checkify.debug_check", "jax._src.checkify.div_checks", "jax._src.checkify.float_checks", "jax._src.checkify.index_checks", "jax._src.checkify.init_error", "jax._src.checkify.nan_checks", "jax._src.checkify.user_checks" ], "n_typable": 41, "n_typed": 17, "n_any": 0, "n_untyped": 24, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 12, "n_methods": 10, "n_method_overloads": 10, "n_method_params": 10, "n_classes": 2, "n_attrs": 13, "n_properties": 0, "n_type_ignores": 1 }, { "path": "jax/experimental/colocated_python/api.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.colocated_python.api.colocated_cpu_devices", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 2, "n_typable": 2, "n_params": 1, "n_function_overloads": 2, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.colocated_python.api.colocated_python", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.colocated_python.api.colocated_python_class", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.colocated_python.api", "names": [ "jax.experimental.colocated_python.api.colocated_cpu_devices", "jax.experimental.colocated_python.api.colocated_python", "jax.experimental.colocated_python.api.colocated_python_class" ], "n_typable": 6, "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_functions": 3, "n_function_overloads": 4, "n_function_params": 3, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/colocated_python/func.py", "symbol_reports": [ { "kind": "class", "name": "jax.experimental.colocated_python.func.FunctionInfo", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "FunctionInfo.fun", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "FunctionInfo.fun_sourceinfo", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "FunctionInfo.fun_signature", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.colocated_python.func.Specialization", "methods": [ { "kind": "function", "name": "Specialization.update", "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Specialization.in_specs_treedef", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Specialization.in_specs_leaves", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Specialization.out_specs_fn", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Specialization.out_specs_treedef", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Specialization.out_specs_leaves", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Specialization.devices", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 7, "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 6, "n_attrs": 6, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.colocated_python.func.make_callable", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax.experimental.colocated_python.func.WeakSpec", "methods": [ { "kind": "function", "name": "WeakSpec.__init__", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "WeakSpec.dtypes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "WeakSpec.shapes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "WeakSpec.sharding_ids", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "WeakSpec.treedef", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 3, "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 2, "n_attrs": 4, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.colocated_python.func.StrongSpec", "methods": [ { "kind": "function", "name": "StrongSpec.__init__", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "StrongSpec.in_specs_treedef", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "StrongSpec.in_specs_leaves", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 3, "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 2, "n_attrs": 2, "n_properties": 0 } ], "type_ignores": [], "name": "jax.experimental.colocated_python.func", "names": [ "jax.experimental.colocated_python.func.FunctionInfo", "jax.experimental.colocated_python.func.Specialization", "jax.experimental.colocated_python.func.StrongSpec", "jax.experimental.colocated_python.func.WeakSpec", "jax.experimental.colocated_python.func.make_callable" ], "n_typable": 17, "n_typed": 13, "n_any": 0, "n_untyped": 4, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 3, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 10, "n_classes": 4, "n_attrs": 15, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/colocated_python/func_backend.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.colocated_python.func_backend.SINGLETON_RESULT_STORE", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [], "name": "jax.experimental.colocated_python.func_backend", "names": [ "jax.experimental.colocated_python.func_backend.SINGLETON_RESULT_STORE" ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/colocated_python/obj.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.colocated_python.obj.wrap_class", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.experimental.colocated_python.obj.SINGLETON_INSTANCE_REGISTRY", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [], "name": "jax.experimental.colocated_python.obj", "names": [ "jax.experimental.colocated_python.obj.SINGLETON_INSTANCE_REGISTRY", "jax.experimental.colocated_python.obj.wrap_class" ], "n_typable": 4, "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 2, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/colocated_python/obj_backend.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.colocated_python.obj_backend.SINGLETON_OBJECT_STORE", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [], "name": "jax.experimental.colocated_python.obj_backend", "names": [ "jax.experimental.colocated_python.obj_backend.SINGLETON_OBJECT_STORE" ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/colocated_python/serialization.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.colocated_python.serialization.cloudpickle", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [ { "kind": "type", "rules": [ "import-not-found" ] } ], "name": "jax.experimental.colocated_python.serialization", "names": [ "jax.experimental.colocated_python.serialization.cloudpickle" ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 1 }, { "path": "jax/_src/compilation_cache.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.compilation_cache.set_cache_dir", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.compilation_cache.reset_cache", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "type_ignores": [], "name": "jax._src.compilation_cache", "names": [ "jax._src.compilation_cache.reset_cache", "jax._src.compilation_cache.set_cache_dir" ], "n_typable": 3, "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 1, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/compute_on.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.compute_on.compute_on", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax._src.compute_on", "names": [ "jax._src.compute_on.compute_on" ], "n_typable": 2, "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 1, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/custom_dce.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.custom_dce.custom_dce_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "class", "name": "jax._src.custom_dce.custom_dce", "methods": [ { "kind": "function", "name": "custom_dce.__init__", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "custom_dce.def_dce", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "custom_dce.__call__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "custom_dce.fun", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "custom_dce.static_argnums", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "custom_dce.dce_rule", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "custom_dce.__getattr__", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 11, "n_typed": 7, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 5, "n_attrs": 4, "n_properties": 0 } ], "type_ignores": [], "name": "jax._src.custom_dce", "names": [ "jax._src.custom_dce.custom_dce", "jax._src.custom_dce.custom_dce_p" ], "n_typable": 12, "n_typed": 7, "n_any": 0, "n_untyped": 5, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 5, "n_classes": 1, "n_attrs": 5, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/custom_partitioning_sharding_rule.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.custom_partitioning_sharding_rule.BATCHING", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "class", "name": "jax._src.custom_partitioning_sharding_rule.CompoundFactor", "methods": [ { "kind": "function", "name": "CompoundFactor.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "CompoundFactor.__new__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [], "n_typable": 4, "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 2, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.custom_partitioning_sharding_rule.ArrayMapping", "methods": [ { "kind": "function", "name": "ArrayMapping.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ArrayMapping.__new__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [], "n_typable": 4, "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 2, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.custom_partitioning_sharding_rule.SdyShardingRule", "methods": [ { "kind": "function", "name": "SdyShardingRule.__init__", "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "SdyShardingRule.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "SdyShardingRule.operand_mappings", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "SdyShardingRule.result_mappings", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "SdyShardingRule.factor_sizes", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "SdyShardingRule.reduction_factors", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "SdyShardingRule.need_replication_factors", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "SdyShardingRule.permutation_factors", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 14, "n_typed": 12, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 6, "n_attrs": 6, "n_properties": 0 } ], "type_ignores": [], "name": "jax._src.custom_partitioning_sharding_rule", "names": [ "jax._src.custom_partitioning_sharding_rule.ArrayMapping", "jax._src.custom_partitioning_sharding_rule.BATCHING", "jax._src.custom_partitioning_sharding_rule.CompoundFactor", "jax._src.custom_partitioning_sharding_rule.SdyShardingRule" ], "n_typable": 23, "n_typed": 13, "n_any": 0, "n_untyped": 10, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 10, "n_classes": 3, "n_attrs": 7, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/custom_partitioning.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.custom_partitioning.custom_partitioning", "methods": [ { "kind": "function", "name": "custom_partitioning.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "custom_partitioning.def_partition", "n_typed": 0, "n_any": 0, "n_untyped": 10, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "custom_partitioning.__call__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "custom_partitioning.fun", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "custom_partitioning.partition", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "custom_partitioning.static_argnums", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "custom_partitioning.propagate_user_sharding", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "custom_partitioning.infer_sharding_from_operands", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "custom_partitioning.sharding_rule", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "custom_partitioning.__getattr__", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 23, "n_typed": 0, "n_any": 1, "n_untyped": 22, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 13, "n_attrs": 7, "n_properties": 0 }, { "kind": "attr", "name": "jax._src.custom_partitioning.custom_partitioning_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [ { "kind": "type", "rules": null } ], "name": "jax._src.custom_partitioning", "names": [ "jax._src.custom_partitioning.custom_partitioning", "jax._src.custom_partitioning.custom_partitioning_p" ], "n_typable": 24, "n_typed": 0, "n_any": 1, "n_untyped": 23, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 13, "n_classes": 1, "n_attrs": 8, "n_properties": 0, "n_type_ignores": 1 }, { "path": "jax/experimental/fused.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.fused.unsafe_zip", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.fused.unsafe_map", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.fused.map", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.fused.zip", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.fused.fused", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.fused.fused_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [], "name": "jax.experimental.fused", "names": [ "jax.experimental.fused.fused", "jax.experimental.fused.fused_p", "jax.experimental.fused.map", "jax.experimental.fused.unsafe_map", "jax.experimental.fused.unsafe_zip", "jax.experimental.fused.zip" ], "n_typable": 3, "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 1, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 5, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/interpreters/ad.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.interpreters.ad.instantiate_zeros", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.interpreters.ad.is_undefined_primal", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax._src.interpreters.ad.custom_lin_p", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax._src.interpreters.ad.deflinear", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.interpreters.ad.linearize", "n_typed": 1, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "class", "name": "jax._src.interpreters.ad.UndefinedPrimal", "methods": [ { "kind": "function", "name": "UndefinedPrimal.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "UndefinedPrimal.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "UndefinedPrimal.aval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 4, "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 1, "n_attrs": 1, "n_properties": 0 }, { "kind": "function", "name": "jax._src.interpreters.ad.add_tangents", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.interpreters.ad.jvp", "n_typed": 1, "n_any": 1, "n_untyped": 3, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.interpreters.ad.defjvp", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.interpreters.ad.JVPTracer", "methods": [ { "kind": "function", "name": "JVPTracer.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "JVPTracer._short_repr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "JVPTracer.cur_qdd", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "JVPTracer.full_lower", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "JVPTracer.to_concrete_value", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "JVPTracer.get_referent", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "JVPTracer.type_state", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "JVPTracer.aval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "JVPTracer.primal", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "JVPTracer.tangent", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 13, "n_typed": 0, "n_any": 0, "n_untyped": 13, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 3, "n_attrs": 2, "n_properties": 1 }, { "kind": "class", "name": "jax._src.interpreters.ad.JVPTrace", "methods": [ { "kind": "function", "name": "JVPTrace.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "JVPTrace.to_primal_tangent_pair", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "JVPTrace.process_primitive", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "JVPTrace.cur_qdd", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "JVPTrace.process_call", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "JVPTrace.process_map", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "JVPTrace.process_custom_jvp_call", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "JVPTrace.process_custom_vjp_call", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "JVPTrace.process_custom_transpose", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "JVPTrace.tag", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "JVPTrace.parent_trace", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "JVPTrace.requires_low", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 43, "n_typed": 0, "n_any": 0, "n_untyped": 43, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 9, "n_method_overloads": 9, "n_method_params": 31, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax._src.interpreters.ad.defjvp2", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax._src.interpreters.ad.primitive_transposes", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax._src.interpreters.ad.deflinear2", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax._src.interpreters.ad.primitive_jvps", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax._src.interpreters.ad.defbilinear", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.interpreters.ad.get_primitive_transpose", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax._src.interpreters.ad", "names": [ "jax._src.interpreters.ad.JVPTrace", "jax._src.interpreters.ad.JVPTracer", "jax._src.interpreters.ad.UndefinedPrimal", "jax._src.interpreters.ad.add_tangents", "jax._src.interpreters.ad.custom_lin_p", "jax._src.interpreters.ad.defbilinear", "jax._src.interpreters.ad.defjvp", "jax._src.interpreters.ad.defjvp2", "jax._src.interpreters.ad.deflinear", "jax._src.interpreters.ad.deflinear2", "jax._src.interpreters.ad.get_primitive_transpose", "jax._src.interpreters.ad.instantiate_zeros", "jax._src.interpreters.ad.is_undefined_primal", "jax._src.interpreters.ad.jvp", "jax._src.interpreters.ad.linearize", "jax._src.interpreters.ad.primitive_jvps", "jax._src.interpreters.ad.primitive_transposes" ], "n_typable": 98, "n_typed": 5, "n_any": 1, "n_untyped": 92, "n_functions": 11, "n_function_overloads": 11, "n_function_params": 24, "n_methods": 18, "n_method_overloads": 18, "n_method_params": 35, "n_classes": 3, "n_attrs": 9, "n_properties": 1, "n_type_ignores": 0 }, { "path": "jax/_src/state/types.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.state.types.AbstractRef", "methods": [ { "kind": "function", "name": "AbstractRef.__init__", "n_typed": 1, "n_any": 2, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "AbstractRef.lo_ty", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AbstractRef.lower_val", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AbstractRef.raise_val", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AbstractRef.update_weak_type", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AbstractRef.update_vma", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AbstractRef.update", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "AbstractRef._len", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AbstractRef.at", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AbstractRef.bitcast", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AbstractRef.reshape", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AbstractRef.transpose", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AbstractRef.T", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AbstractRef.get", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "AbstractRef.swap", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "AbstractRef.set", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "AbstractRef.addupdate", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "AbstractRef._getitem", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "AbstractRef._setitem", "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "AbstractRef._addupdate", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "AbstractRef.str_short", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "AbstractRef.__repr__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AbstractRef.__str__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AbstractRef.to_tangent_aval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AbstractRef.to_cotangent_aval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AbstractRef.__eq__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AbstractRef.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "AbstractRef.is_high", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "AbstractRef.weak_type", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "AbstractRef.shape", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "AbstractRef.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "AbstractRef.sharding", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "AbstractRef.vma", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "AbstractRef.inner_aval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "AbstractRef.memory_space", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "AbstractRef.kind", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "AbstractRef.ndim", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "AbstractRef.size", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 74, "n_typed": 8, "n_any": 2, "n_untyped": 64, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 27, "n_method_overloads": 27, "n_method_params": 36, "n_attrs": 5, "n_properties": 6 }, { "kind": "class", "name": "jax._src.state.types.TransformedRef", "methods": [ { "kind": "function", "name": "TransformedRef.bitcast", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TransformedRef.reshape", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TransformedRef.transpose", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TransformedRef.set", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "TransformedRef.swap", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "TransformedRef.get", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TransformedRef.__getattr__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TransformedRef.__getitem__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TransformedRef.__setitem__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [ { "kind": "property", "name": "TransformedRef.is_dynamic_size", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "TransformedRef.type", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "TransformedRef.shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "TransformedRef.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "TransformedRef.at", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "TransformedRef.ref", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TransformedRef.transforms", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TransformedRef.ndim", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "TransformedRef.size", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "TransformedRef.T", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 29, "n_typed": 4, "n_any": 0, "n_untyped": 25, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 9, "n_method_overloads": 9, "n_method_params": 12, "n_attrs": 5, "n_properties": 5 }, { "kind": "class", "name": "jax._src.state.types.Transform", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.state.types.TransposeTransform", "methods": [ { "kind": "function", "name": "TransposeTransform.undo", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TransposeTransform.transform_type", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TransposeTransform.pretty_print", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "TransposeTransform.permutation", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 6, "n_typed": 4, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 3, "n_attrs": 1, "n_properties": 0 } ], "type_ignores": [ { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null } ], "name": "jax._src.state.types", "names": [ "jax._src.state.types.AbstractRef", "jax._src.state.types.Transform", "jax._src.state.types.TransformedRef", "jax._src.state.types.TransposeTransform" ], "n_typable": 109, "n_typed": 16, "n_any": 2, "n_untyped": 91, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 39, "n_method_overloads": 39, "n_method_params": 51, "n_classes": 4, "n_attrs": 11, "n_properties": 11, "n_type_ignores": 5 }, { "path": "jax/_src/hijax.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.hijax.HiPrimitive", "methods": [ { "kind": "function", "name": "HiPrimitive.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "HiPrimitive.is_high", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "HiPrimitive.is_effectful", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "HiPrimitive.abstract_eval", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "HiPrimitive.to_lojax", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "HiPrimitive.jvp", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "HiPrimitive.transpose", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "HiPrimitive.name", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 21, "n_typed": 2, "n_any": 0, "n_untyped": 19, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 13, "n_attrs": 1, "n_properties": 0 }, { "kind": "class", "name": "jax._src.hijax.MutableHiType", "methods": [ { "kind": "function", "name": "MutableHiType.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "MutableHiType.__eq__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "MutableHiType.lo_ty_qdd", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "MutableHiType.lo_ty", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "MutableHiType.new_from_loval", "n_typed": 1, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "MutableHiType.read_loval", "n_typed": 2, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "MutableHiType.update_from_loval", "n_typed": 2, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "MutableHiType.to_tangent_aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "MutableHiType.to_cotangent_aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "MutableHiType.is_high", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MutableHiType.has_qdd", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MutableHiType.type_state", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 19, "n_typed": 9, "n_any": 5, "n_untyped": 5, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 9, "n_method_overloads": 9, "n_method_params": 9, "n_attrs": 3, "n_properties": 0 }, { "kind": "class", "name": "jax._src.hijax.HipSpec", "methods": [ { "kind": "function", "name": "HipSpec.to_lo", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax._src.hijax.register_hitype", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.hijax.VJPHiPrimitive", "methods": [ { "kind": "function", "name": "VJPHiPrimitive.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "VJPHiPrimitive.expand", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "VJPHiPrimitive.vjp_fwd", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "VJPHiPrimitive.vjp_bwd", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "VJPHiPrimitive.vjp_bwd_retval", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "VJPHiPrimitive.jvp", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "VJPHiPrimitive.lin", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "VJPHiPrimitive.linearized", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "VJPHiPrimitive.batch", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "VJPHiPrimitive.batch_dim_rule", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "VJPHiPrimitive.dce", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "VJPHiPrimitive.remat", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "VJPHiPrimitive.__call__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "VJPHiPrimitive.check", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "VJPHiPrimitive.staging", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "VJPHiPrimitive.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "VJPHiPrimitive.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "VJPHiPrimitive.__eq__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "VJPHiPrimitive.in_avals", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "VJPHiPrimitive.out_aval", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "VJPHiPrimitive.params", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 49, "n_typed": 2, "n_any": 1, "n_untyped": 46, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 18, "n_method_overloads": 18, "n_method_params": 28, "n_attrs": 3, "n_properties": 0 }, { "kind": "class", "name": "jax._src.hijax.HiType", "methods": [ { "kind": "function", "name": "HiType.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "HiType.__eq__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "HiType.lo_ty", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "HiType.lower_val", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "HiType.raise_val", "n_typed": 0, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "HiType.to_tangent_aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "HiType.to_cotangent_aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "HiType.vspace_zero", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "HiType.vspace_add", "n_typed": 0, "n_any": 3, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "HiType.dec_rank", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "HiType.inc_rank", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "HiType.shard", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "HiType.unshard", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "HiType.is_high", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "HiType.has_qdd", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 29, "n_typed": 17, "n_any": 7, "n_untyped": 5, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 13, "n_method_overloads": 13, "n_method_params": 16, "n_attrs": 2, "n_properties": 0 } ], "type_ignores": [ { "kind": "type", "rules": null }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "type", "rules": null }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null } ], "name": "jax._src.hijax", "names": [ "jax._src.hijax.HiPrimitive", "jax._src.hijax.HiType", "jax._src.hijax.HipSpec", "jax._src.hijax.MutableHiType", "jax._src.hijax.VJPHiPrimitive", "jax._src.hijax.register_hitype" ], "n_typable": 122, "n_typed": 31, "n_any": 13, "n_untyped": 78, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 2, "n_methods": 48, "n_method_overloads": 48, "n_method_params": 66, "n_classes": 5, "n_attrs": 9, "n_properties": 0, "n_type_ignores": 23 }, { "path": "jax/_src/export/shape_poly.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.export.shape_poly.PolyShape", "methods": [ { "kind": "function", "name": "PolyShape.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "PolyShape.__new__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "PolyShape.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [], "n_typable": 5, "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 2, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax._src.export.shape_poly.symbolic_shape", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "class", "name": "jax._src.export.shape_poly.SymbolicScope", "methods": [ { "kind": "function", "name": "SymbolicScope.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SymbolicScope.__str__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "SymbolicScope.__repr__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "SymbolicScope._parse_and_process_explicit_constraint", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SymbolicScope._process_explicit_constraint", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SymbolicScope._check_same_scope", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "SymbolicScope._clear_caches", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [], "n_typable": 14, "n_typed": 9, "n_any": 0, "n_untyped": 5, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 7, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax._src.export.shape_poly.is_symbolic_dim", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.export.shape_poly.symbolic_args_specs", "n_typed": 2, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [ { "kind": "type", "rules": [ "bad-return-type", "unused-ignore" ] }, { "kind": "type", "rules": [ "bad-return-type", "unused-ignore" ] }, { "kind": "type", "rules": null } ], "name": "jax._src.export.shape_poly", "names": [ "jax._src.export.shape_poly.PolyShape", "jax._src.export.shape_poly.SymbolicScope", "jax._src.export.shape_poly.is_symbolic_dim", "jax._src.export.shape_poly.symbolic_args_specs", "jax._src.export.shape_poly.symbolic_shape" ], "n_typable": 31, "n_typed": 18, "n_any": 0, "n_untyped": 13, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 9, "n_methods": 10, "n_method_overloads": 10, "n_method_params": 9, "n_classes": 2, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 3 }, { "path": "jax/experimental/jax2tf/jax2tf.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.jax2tf.jax2tf.split_to_logical_devices", "n_typed": 1, "n_any": 1, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.jax2tf.jax2tf.dtype_of_val", "n_typed": 0, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.jax2tf.jax2tf.eval_polymorphic_shape", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.jax2tf.jax2tf.convert", "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.experimental.jax2tf.jax2tf.preprocess_arg_tf", "n_typed": 1, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.jax2tf.jax2tf.inside_call_tf", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "class", "name": "jax.experimental.jax2tf.jax2tf.NativeSerializationImpl", "methods": [ { "kind": "function", "name": "NativeSerializationImpl.__init__", "n_typed": 2, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "NativeSerializationImpl.before_conversion", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "NativeSerializationImpl.after_conversion", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "NativeSerializationImpl.run_fun_tf", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "NativeSerializationImpl.get_vjp_fun", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "NativeSerializationImpl.convert_kwargs", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "NativeSerializationImpl.fun_jax", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "NativeSerializationImpl.args_specs", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "NativeSerializationImpl.kwargs_specs", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "NativeSerializationImpl.native_serialization_disabled_checks", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "NativeSerializationImpl.native_serialization_platforms", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 17, "n_typed": 5, "n_any": 0, "n_untyped": 12, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 6, "n_attrs": 6, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.jax2tf.jax2tf.DEFAULT_NATIVE_SERIALIZATION", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.jax2tf.jax2tf.PrecisionType", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.jax2tf.jax2tf.PartitionsOrReplicated", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.jax2tf.jax2tf.get_thread_local_state_call_tf_concrete_function_list", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "type_ignores": [], "name": "jax.experimental.jax2tf.jax2tf", "names": [ "jax.experimental.jax2tf.jax2tf.DEFAULT_NATIVE_SERIALIZATION", "jax.experimental.jax2tf.jax2tf.NativeSerializationImpl", "jax.experimental.jax2tf.jax2tf.PartitionsOrReplicated", "jax.experimental.jax2tf.jax2tf.PrecisionType", "jax.experimental.jax2tf.jax2tf.convert", "jax.experimental.jax2tf.jax2tf.dtype_of_val", "jax.experimental.jax2tf.jax2tf.eval_polymorphic_shape", "jax.experimental.jax2tf.jax2tf.get_thread_local_state_call_tf_concrete_function_list", "jax.experimental.jax2tf.jax2tf.inside_call_tf", "jax.experimental.jax2tf.jax2tf.preprocess_arg_tf", "jax.experimental.jax2tf.jax2tf.split_to_logical_devices" ], "n_typable": 41, "n_typed": 20, "n_any": 5, "n_untyped": 16, "n_functions": 7, "n_function_overloads": 7, "n_function_params": 15, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 6, "n_classes": 1, "n_attrs": 9, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/jax2tf/call_tf.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.jax2tf.call_tf.call_tf", "n_typed": 2, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "class", "name": "jax.experimental.jax2tf.call_tf.CallTfOrderedEffect", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "CallTfOrderedEffect.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 1, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.jax2tf.call_tf.call_tf_ordered_effect", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "class", "name": "jax.experimental.jax2tf.call_tf.UnspecifiedOutputShapeDtype", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.jax2tf.call_tf.emit_tf_embedded_graph_custom_call", "n_typed": 2, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "attr", "name": "jax.experimental.jax2tf.call_tf.call_tf_effect", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "class", "name": "jax.experimental.jax2tf.call_tf.CallTfEffect", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "CallTfEffect.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 1, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.jax2tf.call_tf.call_tf_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.jax2tf.call_tf.add_to_call_tf_concrete_function_list", "n_typed": 2, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.jax2tf.call_tf.check_tf_result", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax.experimental.jax2tf.call_tf", "names": [ "jax.experimental.jax2tf.call_tf.CallTfEffect", "jax.experimental.jax2tf.call_tf.CallTfOrderedEffect", "jax.experimental.jax2tf.call_tf.UnspecifiedOutputShapeDtype", "jax.experimental.jax2tf.call_tf.add_to_call_tf_concrete_function_list", "jax.experimental.jax2tf.call_tf.call_tf", "jax.experimental.jax2tf.call_tf.call_tf_effect", "jax.experimental.jax2tf.call_tf.call_tf_ordered_effect", "jax.experimental.jax2tf.call_tf.call_tf_p", "jax.experimental.jax2tf.call_tf.check_tf_result", "jax.experimental.jax2tf.call_tf.emit_tf_embedded_graph_custom_call" ], "n_typable": 23, "n_typed": 10, "n_any": 1, "n_untyped": 12, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 16, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 3, "n_attrs": 5, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/jet.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.jet.deflinear", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax.experimental.jet.ZeroSeries", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.jet.zero_series", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.jet.fact", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.jet.zero_prop", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.jet.deriv_prop", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "class", "name": "jax.experimental.jet.ZeroTerm", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.jet.jet_subtrace", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "attr", "name": "jax.experimental.jet.call_param_updaters", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "class", "name": "jax.experimental.jet.JetTracer", "methods": [ { "kind": "function", "name": "JetTracer.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "JetTracer.full_lower", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "JetTracer.aval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "JetTracer.primal", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "JetTracer.terms", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 8, "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 3, "n_attrs": 2, "n_properties": 1 }, { "kind": "attr", "name": "jax.experimental.jet.jet_rules", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "class", "name": "jax.experimental.jet.JetTrace", "methods": [ { "kind": "function", "name": "JetTrace.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "JetTrace.to_primal_terms_pair", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "JetTrace.process_primitive", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "JetTrace.process_call", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "JetTrace.process_custom_jvp_call", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "JetTrace.process_custom_vjp_call", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "JetTrace.tag", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "JetTrace.parent_trace", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "JetTrace.order", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 31, "n_typed": 0, "n_any": 0, "n_untyped": 31, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 22, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.jet.traceable", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.jet.linear_prop", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.jet.def_deriv", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.experimental.jet.jet2", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.jet.defzero", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.jet.jet_fun", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.jet.jet", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "attr", "name": "jax.experimental.jet.zero_term", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.jet.def_comp", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax.experimental.jet", "names": [ "jax.experimental.jet.JetTrace", "jax.experimental.jet.JetTracer", "jax.experimental.jet.ZeroSeries", "jax.experimental.jet.ZeroTerm", "jax.experimental.jet.call_param_updaters", "jax.experimental.jet.def_comp", "jax.experimental.jet.def_deriv", "jax.experimental.jet.deflinear", "jax.experimental.jet.defzero", "jax.experimental.jet.deriv_prop", "jax.experimental.jet.fact", "jax.experimental.jet.jet", "jax.experimental.jet.jet2", "jax.experimental.jet.jet_fun", "jax.experimental.jet.jet_rules", "jax.experimental.jet.jet_subtrace", "jax.experimental.jet.linear_prop", "jax.experimental.jet.traceable", "jax.experimental.jet.zero_prop", "jax.experimental.jet.zero_series", "jax.experimental.jet.zero_term" ], "n_typable": 93, "n_typed": 1, "n_any": 0, "n_untyped": 92, "n_functions": 12, "n_function_overloads": 12, "n_function_params": 38, "n_methods": 8, "n_method_overloads": 8, "n_method_params": 25, "n_classes": 4, "n_attrs": 10, "n_properties": 1, "n_type_ignores": 0 }, { "path": "jax/_src/layout.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.layout.Format", "methods": [ { "kind": "function", "name": "Format.__init__", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Format.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Format.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Format.__eq__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Format.layout", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Format.sharding", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 9, "n_typed": 2, "n_any": 0, "n_untyped": 7, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 3, "n_attrs": 2, "n_properties": 0 }, { "kind": "class", "name": "jax._src.layout.Layout", "methods": [ { "kind": "function", "name": "Layout.__init__", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Layout.from_pjrt_layout", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Layout.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Layout.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Layout.__eq__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Layout.update", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Layout._to_xla_layout", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Layout.check_compatible_aval", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Layout.major_to_minor", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Layout.tiling", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Layout.sub_byte_element_size_in_bits", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Layout.AUTO", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 20, "n_typed": 9, "n_any": 0, "n_untyped": 11, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 8, "n_method_overloads": 8, "n_method_params": 8, "n_attrs": 4, "n_properties": 0 } ], "type_ignores": [], "name": "jax._src.layout", "names": [ "jax._src.layout.Format", "jax._src.layout.Layout" ], "n_typable": 29, "n_typed": 11, "n_any": 0, "n_untyped": 18, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 12, "n_method_overloads": 12, "n_method_params": 11, "n_classes": 2, "n_attrs": 6, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/mesh_utils.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.mesh_utils.create_hybrid_device_mesh", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "attr", "name": "jax._src.mesh_utils.device_kind_handler_dict", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax._src.mesh_utils.create_device_mesh", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax._src.mesh_utils", "names": [ "jax._src.mesh_utils.create_device_mesh", "jax._src.mesh_utils.create_hybrid_device_mesh", "jax._src.mesh_utils.device_kind_handler_dict" ], "n_typable": 13, "n_typed": 13, "n_any": 0, "n_untyped": 0, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 10, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/tpu_custom_call.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.tpu_custom_call.as_tpu_kernel", "n_typed": 21, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 22, "n_params": 21, "n_function_overloads": 1, "n_function_params": 21 }, { "kind": "class", "name": "jax._src.tpu_custom_call.Tiling", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "Tiling.COMPACT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Tiling.SPARSE_CORE", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax._src.tpu_custom_call.lower_module_to_custom_call", "n_typed": 24, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 25, "n_params": 24, "n_function_overloads": 1, "n_function_params": 24 } ], "type_ignores": [ { "kind": "pyrefly", "rules": [ "bad-class-definition" ] } ], "name": "jax._src.tpu_custom_call", "names": [ "jax._src.tpu_custom_call.Tiling", "jax._src.tpu_custom_call.as_tpu_kernel", "jax._src.tpu_custom_call.lower_module_to_custom_call" ], "n_typable": 47, "n_typed": 45, "n_any": 2, "n_untyped": 0, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 45, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 2, "n_properties": 0, "n_type_ignores": 1 }, { "path": "jax/experimental/mosaic/dialects.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.mosaic.dialects.tpu", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [], "name": "jax.experimental.mosaic.dialects", "names": [ "jax.experimental.mosaic.dialects.tpu" ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/mosaic/gpu/fragmented_array.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.WGMMA_COL_LAYOUT", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.TiledLayout", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.WGMMA_ROW_LAYOUT", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.WGMMA_LAYOUT", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.TCGEN05_TRANSPOSED_LAYOUT", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.fragmented_array.optimization_barrier", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 2, "n_typable": 5, "n_params": 4, "n_function_overloads": 2, "n_function_params": 4 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.WGMMA_LAYOUT_UPCAST_4X", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.fragmented_array.tmem_native_layout", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.FragmentedLayout", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.TCGEN05_COL_LAYOUT", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.WGMMA_LAYOUT_UPCAST_2X", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.TCGEN05_LAYOUT", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.fragmented_array.FragmentedArray", "methods": [ { "kind": "function", "name": "FragmentedArray.__init__", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "FragmentedArray.load_strided", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "FragmentedArray.splat", "n_typed": 2, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "FragmentedArray.broadcasted_iota", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "FragmentedArray.to_layout", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray._pointwise", "n_typed": 3, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "FragmentedArray.__pos__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "FragmentedArray.__neg__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "FragmentedArray.__add__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__radd__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__mul__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__rmul__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__sub__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__rsub__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__truediv__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__rtruediv__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__floordiv__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__rfloordiv__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__mod__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__rmod__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__invert__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "FragmentedArray.__or__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__ror__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__and__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__rand__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__xor__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__rxor__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__lshift__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__rshift__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__eq__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__ne__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__lt__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__le__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__gt__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__ge__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray._compare", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "FragmentedArray.max", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.min", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.copysign", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.exp", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.exp2", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.log", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.log2", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.sin", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.cos", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.tanh", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.rsqrt", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.abs", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "FragmentedArray.round", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "FragmentedArray.round_even", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "FragmentedArray.erf", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "FragmentedArray.atan2", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray._lift_fast_instr", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray._lift_fast_packed_instr", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "FragmentedArray.bitcast", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "FragmentedArray.__getitem__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.__setitem__", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "FragmentedArray.copy", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "FragmentedArray.astype", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "FragmentedArray.reduce", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "FragmentedArray.broadcast", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.reshape", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.broadcast_minor", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.broadcast_in_dim", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "FragmentedArray.select", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "FragmentedArray.build", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "FragmentedArray.foreach", "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "FragmentedArray.debug_print", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.store_untiled", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "FragmentedArray.load_reduce_untiled", "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "FragmentedArray.load_untiled", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "FragmentedArray._store_untiled_splat", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FragmentedArray.store_tiled_async", "n_typed": 7, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "FragmentedArray.store_tiled", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "FragmentedArray.load_tiled", "n_typed": 8, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "FragmentedArray.transfer_strided", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "FragmentedArray.transfer_tiled", "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "FragmentedArray.tree_flatten", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "FragmentedArray.tree_unflatten", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [ { "kind": "property", "name": "FragmentedArray.shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "FragmentedArray.mlir_dtype", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "FragmentedArray.registers", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "FragmentedArray.layout", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "FragmentedArray.is_signed", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 216, "n_typed": 125, "n_any": 0, "n_untyped": 91, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 79, "n_method_overloads": 79, "n_method_params": 135, "n_attrs": 3, "n_properties": 2 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.TCGEN05_ROW_LAYOUT", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.fragmented_array.copy_tiled", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.TMA_GATHER_INDICES_LAYOUT", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.fragmented_array.WGStridedFragLayout", "methods": [ { "kind": "function", "name": "WGStridedFragLayout.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "WGStridedFragLayout.from_shaped_type", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "WGStridedFragLayout.registers_element_type", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "WGStridedFragLayout.registers_shape", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "WGStridedFragLayout.shape_from_registers_shape", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "WGStridedFragLayout.thread_idxs", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "WGStridedFragLayout.linear_thread_idxs", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "WGStridedFragLayout.shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "WGStridedFragLayout.vec_size", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 12, "n_typed": 8, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 5, "n_attrs": 2, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.fragmented_array.WGSplatFragLayout", "methods": [ { "kind": "function", "name": "WGSplatFragLayout.can_broadcast_to", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "WGSplatFragLayout.registers_element_type", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "WGSplatFragLayout.registers_shape", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "WGSplatFragLayout.shape_from_registers_shape", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "WGSplatFragLayout.thread_idxs", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "WGSplatFragLayout.shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 10, "n_typed": 7, "n_any": 0, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 5, "n_attrs": 1, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.WGMMA_LAYOUT_8BIT", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.TMEM_NATIVE_LAYOUT", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.WGMMA_TRANSPOSED_LAYOUT", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.fragmented_array.can_relayout_wgmma_2x_to_wgmma", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.fragmented_array.plan_tiled_transfer", "n_typed": 10, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.fragmented_array.subf", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.TrivialTransferPlan", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.SMEM_BANKS", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.fragmented_array.TransferPlan", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.fragmented_array.can_relayout_wgmma_4x_to_wgmma_2x", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.SMEM_BANK_BYTES", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.WARPS_IN_WARPGROUP", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.fragmented_array.addf", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.fragmented_array.StaggeredTransferPlanImpl", "methods": [ { "kind": "function", "name": "StaggeredTransferPlanImpl.select", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "StaggeredTransferPlanImpl.select_if_group", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [ { "kind": "property", "name": "StaggeredTransferPlanImpl.tile_index_transforms", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "StaggeredTransferPlanImpl.stagger", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "StaggeredTransferPlanImpl.dim", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "StaggeredTransferPlanImpl.size", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "StaggeredTransferPlanImpl.group_pred", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 7, "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 4, "n_attrs": 4, "n_properties": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.WGMMA_LAYOUT_ACC_32BIT", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.StaggeredTransferPlan", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.T", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.fragmented_array.TiledLayoutImpl", "methods": [ { "kind": "function", "name": "TiledLayoutImpl.__post_init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TiledLayoutImpl.thread_idxs", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TiledLayoutImpl.registers_element_type", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TiledLayoutImpl.registers_shape", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TiledLayoutImpl.shape_from_registers_shape", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TiledLayoutImpl._delinearize_index", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "TiledLayoutImpl.lane_indices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "TiledLayoutImpl.warp_indices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "TiledLayoutImpl.remove_dimension", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TiledLayoutImpl.reduce", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TiledLayoutImpl.canonicalize", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "TiledLayoutImpl.partitioned_warp_dims", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "TiledLayoutImpl.partitioned_lane_dims", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "TiledLayoutImpl.base_tile_shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "TiledLayoutImpl.tiled_tiling_shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "TiledLayoutImpl.tiled_tiling_rank", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "TiledLayoutImpl.vector_length", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "TiledLayoutImpl.tiling", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TiledLayoutImpl.warp_dims", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TiledLayoutImpl.lane_dims", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TiledLayoutImpl.vector_dim", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TiledLayoutImpl._check_canonical", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 26, "n_typed": 25, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 11, "n_method_overloads": 11, "n_method_params": 9, "n_attrs": 5, "n_properties": 6 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.Tiling", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.WARP_SIZE", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.Replicated", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.fragmented_array.IndexTransform", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.fragmented_array.tiled_copy_smem_gmem_layout", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.fragmented_array.mulf", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.fragmented_array.enumerate_negative", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.fragmented_array.TrivialTransferPlanImpl", "methods": [ { "kind": "function", "name": "TrivialTransferPlanImpl.select", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TrivialTransferPlanImpl.select_if_group", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [ { "kind": "property", "name": "TrivialTransferPlanImpl.tile_index_transforms", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [], "n_typable": 7, "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 4, "n_attrs": 0, "n_properties": 1 } ], "type_ignores": [], "name": "jax.experimental.mosaic.gpu.fragmented_array", "names": [ "jax.experimental.mosaic.gpu.fragmented_array.FragmentedArray", "jax.experimental.mosaic.gpu.fragmented_array.FragmentedLayout", "jax.experimental.mosaic.gpu.fragmented_array.IndexTransform", "jax.experimental.mosaic.gpu.fragmented_array.Replicated", "jax.experimental.mosaic.gpu.fragmented_array.SMEM_BANKS", "jax.experimental.mosaic.gpu.fragmented_array.SMEM_BANK_BYTES", "jax.experimental.mosaic.gpu.fragmented_array.StaggeredTransferPlan", "jax.experimental.mosaic.gpu.fragmented_array.StaggeredTransferPlanImpl", "jax.experimental.mosaic.gpu.fragmented_array.T", "jax.experimental.mosaic.gpu.fragmented_array.TCGEN05_COL_LAYOUT", "jax.experimental.mosaic.gpu.fragmented_array.TCGEN05_LAYOUT", "jax.experimental.mosaic.gpu.fragmented_array.TCGEN05_ROW_LAYOUT", "jax.experimental.mosaic.gpu.fragmented_array.TCGEN05_TRANSPOSED_LAYOUT", "jax.experimental.mosaic.gpu.fragmented_array.TMA_GATHER_INDICES_LAYOUT", "jax.experimental.mosaic.gpu.fragmented_array.TMEM_NATIVE_LAYOUT", "jax.experimental.mosaic.gpu.fragmented_array.TiledLayout", "jax.experimental.mosaic.gpu.fragmented_array.TiledLayoutImpl", "jax.experimental.mosaic.gpu.fragmented_array.Tiling", "jax.experimental.mosaic.gpu.fragmented_array.TransferPlan", "jax.experimental.mosaic.gpu.fragmented_array.TrivialTransferPlan", "jax.experimental.mosaic.gpu.fragmented_array.TrivialTransferPlanImpl", "jax.experimental.mosaic.gpu.fragmented_array.WARPS_IN_WARPGROUP", "jax.experimental.mosaic.gpu.fragmented_array.WARP_SIZE", "jax.experimental.mosaic.gpu.fragmented_array.WGMMA_COL_LAYOUT", "jax.experimental.mosaic.gpu.fragmented_array.WGMMA_LAYOUT", "jax.experimental.mosaic.gpu.fragmented_array.WGMMA_LAYOUT_8BIT", "jax.experimental.mosaic.gpu.fragmented_array.WGMMA_LAYOUT_ACC_32BIT", "jax.experimental.mosaic.gpu.fragmented_array.WGMMA_LAYOUT_UPCAST_2X", "jax.experimental.mosaic.gpu.fragmented_array.WGMMA_LAYOUT_UPCAST_4X", "jax.experimental.mosaic.gpu.fragmented_array.WGMMA_ROW_LAYOUT", "jax.experimental.mosaic.gpu.fragmented_array.WGMMA_TRANSPOSED_LAYOUT", "jax.experimental.mosaic.gpu.fragmented_array.WGSplatFragLayout", "jax.experimental.mosaic.gpu.fragmented_array.WGStridedFragLayout", "jax.experimental.mosaic.gpu.fragmented_array.addf", "jax.experimental.mosaic.gpu.fragmented_array.can_relayout_wgmma_2x_to_wgmma", "jax.experimental.mosaic.gpu.fragmented_array.can_relayout_wgmma_4x_to_wgmma_2x", "jax.experimental.mosaic.gpu.fragmented_array.copy_tiled", "jax.experimental.mosaic.gpu.fragmented_array.enumerate_negative", "jax.experimental.mosaic.gpu.fragmented_array.mulf", "jax.experimental.mosaic.gpu.fragmented_array.optimization_barrier", "jax.experimental.mosaic.gpu.fragmented_array.plan_tiled_transfer", "jax.experimental.mosaic.gpu.fragmented_array.subf", "jax.experimental.mosaic.gpu.fragmented_array.tiled_copy_smem_gmem_layout", "jax.experimental.mosaic.gpu.fragmented_array.tmem_native_layout" ], "n_typable": 340, "n_typed": 217, "n_any": 3, "n_untyped": 120, "n_functions": 11, "n_function_overloads": 12, "n_function_params": 30, "n_methods": 106, "n_method_overloads": 106, "n_method_params": 162, "n_classes": 7, "n_attrs": 41, "n_properties": 10, "n_type_ignores": 0 }, { "path": "jax/experimental/mosaic/gpu/core.py", "symbol_reports": [ { "kind": "class", "name": "jax.experimental.mosaic.gpu.core.Barrier", "methods": [ { "kind": "function", "name": "Barrier.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Barrier.arrival_count", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Barrier.num_barriers", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.core.ClusterBarrier", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "ClusterBarrier.collective_dims", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ClusterBarrier.arrival_count", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ClusterBarrier.num_barriers", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.core.Union", "methods": [ { "kind": "function", "name": "Union.__iter__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Union.members", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 1, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.core.TMEM", "methods": [ { "kind": "function", "name": "TMEM.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "TMEM.shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TMEM.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TMEM._", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TMEM.layout", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TMEM.collective", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TMEM.packing", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 6, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.core.TMABarrier", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "TMABarrier.num_barriers", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 1, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.core.supports_cross_device_collectives", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.core.as_gpu_kernel", "n_typed": 9, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 14, "n_params": 13, "n_function_overloads": 1, "n_function_params": 13 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.core.LoweringSemantics", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "LoweringSemantics.Lane", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringSemantics.Warpgroup", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.core.as_torch_gpu_kernel", "n_typed": 8, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 13, "n_params": 12, "n_function_overloads": 1, "n_function_params": 12 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.core.KNOWN_KERNELS", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.core.is_nvshmem_available", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.core.cuda_root", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.core.mosaic_gpu_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.core.is_single_process_multi_device_topology", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.core.libdevice_path", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.core.T", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.core.RUNTIME_PATH", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.core.FWD_COMPAT_IR_VERSION", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.core.PYTHON_RUNFILES", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.core.artificial_shared_memory_limit", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.mosaic.gpu.core", "names": [ "jax.experimental.mosaic.gpu.core.Barrier", "jax.experimental.mosaic.gpu.core.ClusterBarrier", "jax.experimental.mosaic.gpu.core.FWD_COMPAT_IR_VERSION", "jax.experimental.mosaic.gpu.core.KNOWN_KERNELS", "jax.experimental.mosaic.gpu.core.LoweringSemantics", "jax.experimental.mosaic.gpu.core.PYTHON_RUNFILES", "jax.experimental.mosaic.gpu.core.RUNTIME_PATH", "jax.experimental.mosaic.gpu.core.T", "jax.experimental.mosaic.gpu.core.TMABarrier", "jax.experimental.mosaic.gpu.core.TMEM", "jax.experimental.mosaic.gpu.core.Union", "jax.experimental.mosaic.gpu.core.artificial_shared_memory_limit", "jax.experimental.mosaic.gpu.core.as_gpu_kernel", "jax.experimental.mosaic.gpu.core.as_torch_gpu_kernel", "jax.experimental.mosaic.gpu.core.cuda_root", "jax.experimental.mosaic.gpu.core.is_nvshmem_available", "jax.experimental.mosaic.gpu.core.is_single_process_multi_device_topology", "jax.experimental.mosaic.gpu.core.libdevice_path", "jax.experimental.mosaic.gpu.core.mosaic_gpu_p", "jax.experimental.mosaic.gpu.core.supports_cross_device_collectives" ], "n_typable": 39, "n_typed": 18, "n_any": 0, "n_untyped": 21, "n_functions": 6, "n_function_overloads": 6, "n_function_params": 26, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 0, "n_classes": 6, "n_attrs": 23, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/mosaic/gpu/launch_context.py", "symbol_reports": [ { "kind": "class", "name": "jax.experimental.mosaic.gpu.launch_context.AsyncCopyImplementation", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "AsyncCopyImplementation.TMA", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "AsyncCopyImplementation.CP_ASYNC", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.launch_context.GLOBAL_BROADCAST", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.launch_context.Rounding", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "Rounding.UP", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Rounding.DOWN", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.launch_context.LaunchContext", "methods": [ { "kind": "function", "name": "LaunchContext.named_region", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "LaunchContext._alloc_scratch", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "LaunchContext._recompute_peer_id", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "LaunchContext._get_tma_desc", "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "LaunchContext._prepare_async_copy", "n_typed": 5, "n_any": 1, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "LaunchContext._prepare_tma", "n_typed": 9, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 11, "n_params": 10, "n_function_overloads": 1, "n_function_params": 10 }, { "kind": "function", "name": "LaunchContext.async_copy", "n_typed": 12, "n_any": 1, "n_untyped": 1, "n_overloads": 1, "n_typable": 14, "n_params": 13, "n_function_overloads": 1, "n_function_params": 13 }, { "kind": "function", "name": "LaunchContext.async_prefetch", "n_typed": 7, "n_any": 1, "n_untyped": 1, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "LaunchContext.await_async_copy", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "LaunchContext.await_cp_async_copy", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "LaunchContext._ensure_nvshmem_decls", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "LaunchContext._find_kernel_argument_index", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "LaunchContext.to_remote", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "LaunchContext.to_remote_multicast", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "LaunchContext.device_id", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "LaunchContext.module", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LaunchContext.scratch", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LaunchContext.cluster_size", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LaunchContext.profiler", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LaunchContext.device_collective_metadata", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LaunchContext.host_collective_metadata", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LaunchContext.num_peers", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LaunchContext.tma_descriptors", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LaunchContext.is_device_collective", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 77, "n_typed": 58, "n_any": 3, "n_untyped": 16, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 15, "n_method_overloads": 15, "n_method_params": 62, "n_attrs": 9, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.launch_context.MemRefTransform", "methods": [ { "kind": "function", "name": "MemRefTransform.apply", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "MemRefTransform.transform_index", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "MemRefTransform.transform_shape", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "MemRefTransform.transform_strides", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "MemRefTransform.batch", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "MemRefTransform.to_attr", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [], "n_typable": 11, "n_typed": 11, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 5, "n_attrs": 0, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.launch_context.TMAReductionOp", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.launch_context.SwizzleTransform", "methods": [ { "kind": "function", "name": "SwizzleTransform.to_attr", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "SwizzleTransform.swizzle", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 1, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.launch_context.TileTransform", "methods": [ { "kind": "function", "name": "TileTransform.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "TileTransform.apply", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TileTransform.transform_index", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TileTransform.transform_shape", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TileTransform.transform_strides", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TileTransform.batch", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TileTransform.to_attr", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "TileTransform.tiling", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TileTransform.rounding", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TileTransform._cc_impl", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 12, "n_typed": 11, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 5, "n_attrs": 3, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.launch_context.TransposeTransform", "methods": [ { "kind": "function", "name": "TransposeTransform.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "TransposeTransform.apply", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TransposeTransform.transform_index", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TransposeTransform.transform_shape", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TransposeTransform.transform_strides", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TransposeTransform.batch", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TransposeTransform.to_attr", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "TransposeTransform.permutation", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 12, "n_typed": 11, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 5, "n_attrs": 1, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.launch_context.Scratch", "methods": [ { "kind": "function", "name": "Scratch.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Scratch._find_first_op", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Scratch._create_ops", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Scratch._find_alloc_load_and_device_ptr", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Scratch.device_ptr", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Scratch.finalize_size", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Scratch.next_offset", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Scratch.host_init", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 12, "n_typed": 9, "n_any": 0, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 4, "n_attrs": 2, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.launch_context.ORIGINAL_KERNEL_ARG_ATTR", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.launch_context.MOSAIC_GPU_SMEM_ALLOC_ATTR", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.launch_context.TMA_DESCRIPTOR_BYTES", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.launch_context.DEVICE_ID_ATTR", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.launch_context.GlobalBroadcast", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.launch_context.TMA_DESCRIPTOR_ALIGNMENT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.launch_context.COLLECTIVE_ATTR", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.launch_context.COLLECTIVE_METADATA_SIZE", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.launch_context.CollapseLeadingIndicesTransform", "methods": [ { "kind": "function", "name": "CollapseLeadingIndicesTransform.apply", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "CollapseLeadingIndicesTransform.transform_index", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "CollapseLeadingIndicesTransform.transform_shape", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "CollapseLeadingIndicesTransform.batch", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "CollapseLeadingIndicesTransform.common_stride", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "CollapseLeadingIndicesTransform.strides", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 9, "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 4, "n_attrs": 1, "n_properties": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.launch_context.KERNEL_ARG_ID_ATTR", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.launch_context.ReplicationError", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.launch_context.uses_collective_metadata", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.mosaic.gpu.launch_context", "names": [ "jax.experimental.mosaic.gpu.launch_context.AsyncCopyImplementation", "jax.experimental.mosaic.gpu.launch_context.COLLECTIVE_ATTR", "jax.experimental.mosaic.gpu.launch_context.COLLECTIVE_METADATA_SIZE", "jax.experimental.mosaic.gpu.launch_context.CollapseLeadingIndicesTransform", "jax.experimental.mosaic.gpu.launch_context.DEVICE_ID_ATTR", "jax.experimental.mosaic.gpu.launch_context.GLOBAL_BROADCAST", "jax.experimental.mosaic.gpu.launch_context.GlobalBroadcast", "jax.experimental.mosaic.gpu.launch_context.KERNEL_ARG_ID_ATTR", "jax.experimental.mosaic.gpu.launch_context.LaunchContext", "jax.experimental.mosaic.gpu.launch_context.MOSAIC_GPU_SMEM_ALLOC_ATTR", "jax.experimental.mosaic.gpu.launch_context.MemRefTransform", "jax.experimental.mosaic.gpu.launch_context.ORIGINAL_KERNEL_ARG_ATTR", "jax.experimental.mosaic.gpu.launch_context.ReplicationError", "jax.experimental.mosaic.gpu.launch_context.Rounding", "jax.experimental.mosaic.gpu.launch_context.Scratch", "jax.experimental.mosaic.gpu.launch_context.SwizzleTransform", "jax.experimental.mosaic.gpu.launch_context.TMAReductionOp", "jax.experimental.mosaic.gpu.launch_context.TMA_DESCRIPTOR_ALIGNMENT", "jax.experimental.mosaic.gpu.launch_context.TMA_DESCRIPTOR_BYTES", "jax.experimental.mosaic.gpu.launch_context.TileTransform", "jax.experimental.mosaic.gpu.launch_context.TransposeTransform", "jax.experimental.mosaic.gpu.launch_context.uses_collective_metadata" ], "n_typable": 138, "n_typed": 111, "n_any": 3, "n_untyped": 24, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 1, "n_methods": 46, "n_method_overloads": 46, "n_method_params": 85, "n_classes": 11, "n_attrs": 31, "n_properties": 1, "n_type_ignores": 0 }, { "path": "jax/experimental/mosaic/gpu/utils.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.memref_reshape", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.debug_print", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.bitwidth", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.c", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.memref_fold", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.single_thread", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.when", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.utils.MultimemReductionOp", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.nanosleep", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.utils.DynamicSlice", "methods": [ { "kind": "function", "name": "DynamicSlice.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "DynamicSlice.base", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DynamicSlice.length", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.fori", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.memref_transpose", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.try_cluster_cancel", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.utils.Partition1D", "methods": [ { "kind": "function", "name": "Partition1D.__init__", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "Partition1D.get_base", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Partition1D.refine", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [ { "kind": "property", "name": "Partition1D.num_chunks", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "Partition1D.partition", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Partition1D.base_offset", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 14, "n_typed": 11, "n_any": 0, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 8, "n_attrs": 2, "n_properties": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.is_known_divisible", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.bytewidth", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.get_cluster_ref", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.query_cluster_cancel", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.warpgroup_idx", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.memref_unsqueeze", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.utils.CollectiveBarrierRef", "methods": [ { "kind": "function", "name": "CollectiveBarrierRef.initialize", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "CollectiveBarrierRef.__iter__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "CollectiveBarrierRef.__getitem__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "CollectiveBarrierRef.arrive", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "CollectiveBarrierRef.wait", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "CollectiveBarrierRef.wait_parity", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "CollectiveBarrierRef.barrier", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CollectiveBarrierRef.cluster_mask", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 16, "n_typed": 6, "n_any": 0, "n_untyped": 10, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 10, "n_attrs": 2, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.utils.Partition", "methods": [ { "kind": "function", "name": "Partition.__init__", "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "Partition.get_base", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "Partition.num_chunks", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Partition.target_block_shape", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "Partition.source_bounds", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Partition.target_bounds", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Partition.partition", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Partition.base_offset", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 14, "n_typed": 12, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 6, "n_attrs": 4, "n_properties": 2 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.utils.DialectBarrierRef", "methods": [ { "kind": "function", "name": "DialectBarrierRef.initialize", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "DialectBarrierRef.__iter__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DialectBarrierRef.__getitem__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DialectBarrierRef.wait_parity", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "DialectBarrierRef.wait", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DialectBarrierRef.update_parities", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DialectBarrierRef.arrive", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DialectBarrierRef.arrive_expect_tx", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DialectBarrierRef.get_ptr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DialectBarrierRef.as_barrier_memref", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DialectBarrierRef.from_barrier_memref", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "DialectBarrierRef.barrier_ref", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 21, "n_typed": 13, "n_any": 0, "n_untyped": 8, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 11, "n_method_overloads": 11, "n_method_params": 10, "n_attrs": 1, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.commit_shared", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.utils.ThreadSubset", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "ThreadSubset.WARP", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ThreadSubset.WARPGROUP", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ThreadSubset.BLOCK", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.warpgroup_barrier", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.memref_unfold", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.utils.BarrierRef", "methods": [ { "kind": "function", "name": "BarrierRef.initialize", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "BarrierRef.__iter__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BarrierRef.__getitem__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BarrierRef.wait_parity", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "BarrierRef.wait", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BarrierRef.update_parities", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BarrierRef.arrive", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "BarrierRef.arrive_expect_tx", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "BarrierRef.complete_tx", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "BarrierRef.get_ptr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "BarrierRef.base_address", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BarrierRef.offset", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BarrierRef.phases", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BarrierRef.num_barriers", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 25, "n_typed": 17, "n_any": 0, "n_untyped": 8, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 10, "n_method_overloads": 10, "n_method_params": 15, "n_attrs": 4, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.warp_idx", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.utils.thread_idx", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.memref_slice", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.utils.SemaphoreRef", "methods": [ { "kind": "function", "name": "SemaphoreRef.signal", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "SemaphoreRef.signal_multimem", "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "SemaphoreRef.wait", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "SemaphoreRef.ptr", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 12, "n_typed": 7, "n_any": 0, "n_untyped": 5, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 9, "n_attrs": 1, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.single_thread_predicate", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.tile_shape", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.dtype_to_ir_type", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.memref_ptr", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.utils.MultimemRef", "methods": [ { "kind": "function", "name": "MultimemRef.store", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [ { "kind": "property", "name": "MultimemRef.type", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "MultimemRef.ref", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 4, "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 2, "n_attrs": 1, "n_properties": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.utils.MBARRIER_BYTES", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.getelementptr", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.warp_barrier", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.utils.block_idx", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.redux", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.utils.DYNAMIC", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.warp_tree_reduce", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.get_arch", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.vector_concat", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.utils.WARPS_IN_WARPGROUP", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.utils.WARPGROUP_SIZE", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.prmt", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.smid", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.is_smem_ref", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.get_cluster_ptr", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.vector_slice", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.utils.DYNAMIC32", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.shfl_bfly", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.cluster_idx", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.smem", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.gpu_address_space_to_nvptx", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.multimem_store", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.fence_release_sys", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.is_tmem_ref", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.is_signed", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.get_contiguous_strides", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.parse_indices", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.utils.ForResult", "methods": [], "properties": [ { "kind": "property", "name": "ForResult.result", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "ForResult.op", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ForResult.results", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.ptr_as_memref", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.utils.WARP_SIZE", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.pack_array", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.multimem_load_reduce", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.dyn_dot", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.clock", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.ceil_div", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.bitwidth_impl", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.cluster_collective_mask", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.bitcast", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.utils.Arch", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "Arch.major", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Arch.minor", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.nvvm_mbarrier_arrive_expect_tx", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.is_memref_transposed", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.utils.WORKGROUP_NVPTX_ADDRESS_SPACE", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.globaltimer", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.utils.tmem", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "type_ignores": [ { "kind": "type", "rules": null } ], "name": "jax.experimental.mosaic.gpu.utils", "names": [ "jax.experimental.mosaic.gpu.utils.Arch", "jax.experimental.mosaic.gpu.utils.BarrierRef", "jax.experimental.mosaic.gpu.utils.CollectiveBarrierRef", "jax.experimental.mosaic.gpu.utils.DYNAMIC", "jax.experimental.mosaic.gpu.utils.DYNAMIC32", "jax.experimental.mosaic.gpu.utils.DialectBarrierRef", "jax.experimental.mosaic.gpu.utils.DynamicSlice", "jax.experimental.mosaic.gpu.utils.ForResult", "jax.experimental.mosaic.gpu.utils.MBARRIER_BYTES", "jax.experimental.mosaic.gpu.utils.MultimemReductionOp", "jax.experimental.mosaic.gpu.utils.MultimemRef", "jax.experimental.mosaic.gpu.utils.Partition", "jax.experimental.mosaic.gpu.utils.Partition1D", "jax.experimental.mosaic.gpu.utils.SemaphoreRef", "jax.experimental.mosaic.gpu.utils.ThreadSubset", "jax.experimental.mosaic.gpu.utils.WARPGROUP_SIZE", "jax.experimental.mosaic.gpu.utils.WARPS_IN_WARPGROUP", "jax.experimental.mosaic.gpu.utils.WARP_SIZE", "jax.experimental.mosaic.gpu.utils.WORKGROUP_NVPTX_ADDRESS_SPACE", "jax.experimental.mosaic.gpu.utils.bitcast", "jax.experimental.mosaic.gpu.utils.bitwidth", "jax.experimental.mosaic.gpu.utils.bitwidth_impl", "jax.experimental.mosaic.gpu.utils.block_idx", "jax.experimental.mosaic.gpu.utils.bytewidth", "jax.experimental.mosaic.gpu.utils.c", "jax.experimental.mosaic.gpu.utils.ceil_div", "jax.experimental.mosaic.gpu.utils.clock", "jax.experimental.mosaic.gpu.utils.cluster_collective_mask", "jax.experimental.mosaic.gpu.utils.cluster_idx", "jax.experimental.mosaic.gpu.utils.commit_shared", "jax.experimental.mosaic.gpu.utils.debug_print", "jax.experimental.mosaic.gpu.utils.dtype_to_ir_type", "jax.experimental.mosaic.gpu.utils.dyn_dot", "jax.experimental.mosaic.gpu.utils.fence_release_sys", "jax.experimental.mosaic.gpu.utils.fori", "jax.experimental.mosaic.gpu.utils.get_arch", "jax.experimental.mosaic.gpu.utils.get_cluster_ptr", "jax.experimental.mosaic.gpu.utils.get_cluster_ref", "jax.experimental.mosaic.gpu.utils.get_contiguous_strides", "jax.experimental.mosaic.gpu.utils.getelementptr", "jax.experimental.mosaic.gpu.utils.globaltimer", "jax.experimental.mosaic.gpu.utils.gpu_address_space_to_nvptx", "jax.experimental.mosaic.gpu.utils.is_known_divisible", "jax.experimental.mosaic.gpu.utils.is_memref_transposed", "jax.experimental.mosaic.gpu.utils.is_signed", "jax.experimental.mosaic.gpu.utils.is_smem_ref", "jax.experimental.mosaic.gpu.utils.is_tmem_ref", "jax.experimental.mosaic.gpu.utils.memref_fold", "jax.experimental.mosaic.gpu.utils.memref_ptr", "jax.experimental.mosaic.gpu.utils.memref_reshape", "jax.experimental.mosaic.gpu.utils.memref_slice", "jax.experimental.mosaic.gpu.utils.memref_transpose", "jax.experimental.mosaic.gpu.utils.memref_unfold", "jax.experimental.mosaic.gpu.utils.memref_unsqueeze", "jax.experimental.mosaic.gpu.utils.multimem_load_reduce", "jax.experimental.mosaic.gpu.utils.multimem_store", "jax.experimental.mosaic.gpu.utils.nanosleep", "jax.experimental.mosaic.gpu.utils.nvvm_mbarrier_arrive_expect_tx", "jax.experimental.mosaic.gpu.utils.pack_array", "jax.experimental.mosaic.gpu.utils.parse_indices", "jax.experimental.mosaic.gpu.utils.prmt", "jax.experimental.mosaic.gpu.utils.ptr_as_memref", "jax.experimental.mosaic.gpu.utils.query_cluster_cancel", "jax.experimental.mosaic.gpu.utils.redux", "jax.experimental.mosaic.gpu.utils.shfl_bfly", "jax.experimental.mosaic.gpu.utils.single_thread", "jax.experimental.mosaic.gpu.utils.single_thread_predicate", "jax.experimental.mosaic.gpu.utils.smem", "jax.experimental.mosaic.gpu.utils.smid", "jax.experimental.mosaic.gpu.utils.thread_idx", "jax.experimental.mosaic.gpu.utils.tile_shape", "jax.experimental.mosaic.gpu.utils.tmem", "jax.experimental.mosaic.gpu.utils.try_cluster_cancel", "jax.experimental.mosaic.gpu.utils.vector_concat", "jax.experimental.mosaic.gpu.utils.vector_slice", "jax.experimental.mosaic.gpu.utils.warp_barrier", "jax.experimental.mosaic.gpu.utils.warp_idx", "jax.experimental.mosaic.gpu.utils.warp_tree_reduce", "jax.experimental.mosaic.gpu.utils.warpgroup_barrier", "jax.experimental.mosaic.gpu.utils.warpgroup_idx", "jax.experimental.mosaic.gpu.utils.when" ], "n_typable": 275, "n_typed": 162, "n_any": 0, "n_untyped": 113, "n_functions": 60, "n_function_overloads": 60, "n_function_params": 100, "n_methods": 37, "n_method_overloads": 37, "n_method_params": 60, "n_classes": 11, "n_attrs": 34, "n_properties": 5, "n_type_ignores": 1 }, { "path": "jax/experimental/mosaic/gpu/mma.py", "symbol_reports": [ { "kind": "class", "name": "jax.experimental.mosaic.gpu.mma.MMALayouts", "methods": [ { "kind": "function", "name": "MMALayouts.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "MMALayouts.lhs", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "MMALayouts.rhs", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "MMALayouts.acc", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 5, "n_typed": 1, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 1, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.mma.mma", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.mma.SUPPORTED_F8_TYPES", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [], "name": "jax.experimental.mosaic.gpu.mma", "names": [ "jax.experimental.mosaic.gpu.mma.MMALayouts", "jax.experimental.mosaic.gpu.mma.SUPPORTED_F8_TYPES", "jax.experimental.mosaic.gpu.mma.mma" ], "n_typable": 9, "n_typed": 5, "n_any": 0, "n_untyped": 4, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 3, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 1, "n_classes": 1, "n_attrs": 4, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/mosaic/gpu/__init__.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.mosaic.gpu.ds", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [], "name": "jax.experimental.mosaic.gpu", "names": [ "jax.experimental.mosaic.gpu.ds" ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/mosaic/gpu/layouts.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.mosaic.gpu.layouts.to_layout_attr", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layouts.from_layout_attr", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layouts.to_strided_fragmented_layout_attr", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layouts.splat_is_compatible_with_tiled", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layouts.is_tiled_layout", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layouts.from_transform_attr", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layouts.is_splat_fragmented_layout", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layouts.to_tiled_layout_attr", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layouts.to_transform_attr", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layouts.from_tiled_layout_attr", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layouts.is_tile_transform", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layouts.is_swizzle_transform", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layouts.is_strided_fragmented_layout", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layouts.from_splat_fragmented_layout_attr", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layouts.is_transpose_transform", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layouts.to_splat_fragmented_layout_attr", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layouts.from_strided_fragmented_layout_attr", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.mosaic.gpu.layouts", "names": [ "jax.experimental.mosaic.gpu.layouts.from_layout_attr", "jax.experimental.mosaic.gpu.layouts.from_splat_fragmented_layout_attr", "jax.experimental.mosaic.gpu.layouts.from_strided_fragmented_layout_attr", "jax.experimental.mosaic.gpu.layouts.from_tiled_layout_attr", "jax.experimental.mosaic.gpu.layouts.from_transform_attr", "jax.experimental.mosaic.gpu.layouts.is_splat_fragmented_layout", "jax.experimental.mosaic.gpu.layouts.is_strided_fragmented_layout", "jax.experimental.mosaic.gpu.layouts.is_swizzle_transform", "jax.experimental.mosaic.gpu.layouts.is_tile_transform", "jax.experimental.mosaic.gpu.layouts.is_tiled_layout", "jax.experimental.mosaic.gpu.layouts.is_transpose_transform", "jax.experimental.mosaic.gpu.layouts.splat_is_compatible_with_tiled", "jax.experimental.mosaic.gpu.layouts.to_layout_attr", "jax.experimental.mosaic.gpu.layouts.to_splat_fragmented_layout_attr", "jax.experimental.mosaic.gpu.layouts.to_strided_fragmented_layout_attr", "jax.experimental.mosaic.gpu.layouts.to_tiled_layout_attr", "jax.experimental.mosaic.gpu.layouts.to_transform_attr" ], "n_typable": 35, "n_typed": 35, "n_any": 0, "n_untyped": 0, "n_functions": 17, "n_function_overloads": 17, "n_function_params": 18, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/mosaic/gpu/wgmma.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.mosaic.gpu.wgmma.wgmma", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.wgmma.WGMMAAccumulator", "methods": [ { "kind": "function", "name": "WGMMAAccumulator.__init__", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "WGMMAAccumulator.zero", "n_typed": 1, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "WGMMAAccumulator.from_registers", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "WGMMAAccumulator.tree_flatten", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "WGMMAAccumulator.tree_unflatten", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [ { "kind": "property", "name": "WGMMAAccumulator.value", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "WGMMAAccumulator._original_layout", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "WGMMAAccumulator._value", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 16, "n_typed": 5, "n_any": 0, "n_untyped": 11, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 10, "n_attrs": 2, "n_properties": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.wgmma.wgmma_m64", "n_typed": 9, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 11, "n_params": 10, "n_function_overloads": 1, "n_function_params": 10 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.wgmma.wgmma_fence", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.mosaic.gpu.wgmma", "names": [ "jax.experimental.mosaic.gpu.wgmma.WGMMAAccumulator", "jax.experimental.mosaic.gpu.wgmma.wgmma", "jax.experimental.mosaic.gpu.wgmma.wgmma_fence", "jax.experimental.mosaic.gpu.wgmma.wgmma_m64" ], "n_typable": 34, "n_typed": 20, "n_any": 0, "n_untyped": 14, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 15, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 10, "n_classes": 1, "n_attrs": 2, "n_properties": 1, "n_type_ignores": 0 }, { "path": "jax/experimental/mosaic/gpu/dialect_lowering.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.mosaic.gpu.dialect_lowering.lower_mgpu_dialect", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.dialect_lowering.unwrap_transformed_memref", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.dialect_lowering.MlirLoweringRuleResult", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.dialect_lowering.MlirLoweringRule", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.dialect_lowering.LoweringContext", "methods": [ { "kind": "function", "name": "LoweringContext.check_collective", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "LoweringContext.lower_op", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "LoweringContext.launch_context", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringContext.single_thread_per_block_predicate", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringContext.single_thread_per_warpgroup_predicate", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringContext.single_warp_per_block_predicate", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringContext.auto_barriers", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringContext.smem_requested_bytes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringContext.lowered_operations", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringContext.is_collective_kernel", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 4, "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 2, "n_attrs": 8, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.dialect_lowering.wrap_transformed_memref", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.dialect_lowering.CMPF_IMPLS", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.dialect_lowering.pprint_layout", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.dialect_lowering.swizzle_and_transforms_from_transforms_attr", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.dialect_lowering.Recursed", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.dialect_lowering.tile_strides", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.dialect_lowering.tile_offset", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.dialect_lowering.transform_type", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.dialect_lowering.CMPI_IMPLS", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.dialect_lowering.RECURSED", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.dialect_lowering.fragmented_array_to_ir", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [], "name": "jax.experimental.mosaic.gpu.dialect_lowering", "names": [ "jax.experimental.mosaic.gpu.dialect_lowering.CMPF_IMPLS", "jax.experimental.mosaic.gpu.dialect_lowering.CMPI_IMPLS", "jax.experimental.mosaic.gpu.dialect_lowering.LoweringContext", "jax.experimental.mosaic.gpu.dialect_lowering.MlirLoweringRule", "jax.experimental.mosaic.gpu.dialect_lowering.MlirLoweringRuleResult", "jax.experimental.mosaic.gpu.dialect_lowering.RECURSED", "jax.experimental.mosaic.gpu.dialect_lowering.Recursed", "jax.experimental.mosaic.gpu.dialect_lowering.fragmented_array_to_ir", "jax.experimental.mosaic.gpu.dialect_lowering.lower_mgpu_dialect", "jax.experimental.mosaic.gpu.dialect_lowering.pprint_layout", "jax.experimental.mosaic.gpu.dialect_lowering.swizzle_and_transforms_from_transforms_attr", "jax.experimental.mosaic.gpu.dialect_lowering.tile_offset", "jax.experimental.mosaic.gpu.dialect_lowering.tile_strides", "jax.experimental.mosaic.gpu.dialect_lowering.transform_type", "jax.experimental.mosaic.gpu.dialect_lowering.unwrap_transformed_memref", "jax.experimental.mosaic.gpu.dialect_lowering.wrap_transformed_memref" ], "n_typable": 33, "n_typed": 30, "n_any": 0, "n_untyped": 3, "n_functions": 9, "n_function_overloads": 9, "n_function_params": 18, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 2, "n_classes": 2, "n_attrs": 13, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/mosaic/gpu/layout_inference.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.mosaic.gpu.layout_inference.infer_layout", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.layout_inference.VariableType", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "VariableType.OPERAND", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "VariableType.RESULT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "VariableType.ARGUMENT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layout_inference.vector_value_sites", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layout_inference.is_valid_smem_layout_assignment", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layout_inference.dynamic_gcd", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layout_inference.conjure_assignment", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layout_inference.is_vector", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.layout_inference.ValueSite", "methods": [ { "kind": "function", "name": "ValueSite.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ValueSite.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "ValueSite.value", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "ValueSite.shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "ValueSite.memory_space", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "ValueSite.operation", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ValueSite.type", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ValueSite.index", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ValueSite.region_index", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 5, "n_typed": 3, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 0, "n_attrs": 4, "n_properties": 3 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layout_inference.find_assignments_for", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layout_inference.derive_relayout_constraints", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layout_inference.is_terminator", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.layout_inference.ConstraintSystemDerivationRuleResult", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layout_inference.prime_decomposition", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.layout_inference.MemorySpace", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "MemorySpace.REG", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MemorySpace.SMEM", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MemorySpace.TMEM", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layout_inference.check_layout_assignment", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.layout_inference.DerivationContext", "methods": [ { "kind": "function", "name": "DerivationContext.update", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "DerivationContext.producer_ref", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "DerivationContext.variable_for_value_site", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DerivationContext.value_sites_for_variable", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 4, "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 2, "n_attrs": 2, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.layout_inference.ValueSitesForVariable", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layout_inference.extract_assignment_candidates_from_reduce_equation", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layout_inference.is_valid_register_layout_assignment", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layout_inference.producer_result", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layout_inference.is_valid_tmem_layout_assignment", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.layout_inference.ConstraintSystemDerivationRule", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layout_inference.traverse_op", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layout_inference.consumer_operands", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.layout_inference.assign_layouts", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.mosaic.gpu.layout_inference", "names": [ "jax.experimental.mosaic.gpu.layout_inference.ConstraintSystemDerivationRule", "jax.experimental.mosaic.gpu.layout_inference.ConstraintSystemDerivationRuleResult", "jax.experimental.mosaic.gpu.layout_inference.DerivationContext", "jax.experimental.mosaic.gpu.layout_inference.MemorySpace", "jax.experimental.mosaic.gpu.layout_inference.ValueSite", "jax.experimental.mosaic.gpu.layout_inference.ValueSitesForVariable", "jax.experimental.mosaic.gpu.layout_inference.VariableType", "jax.experimental.mosaic.gpu.layout_inference.assign_layouts", "jax.experimental.mosaic.gpu.layout_inference.check_layout_assignment", "jax.experimental.mosaic.gpu.layout_inference.conjure_assignment", "jax.experimental.mosaic.gpu.layout_inference.consumer_operands", "jax.experimental.mosaic.gpu.layout_inference.derive_relayout_constraints", "jax.experimental.mosaic.gpu.layout_inference.dynamic_gcd", "jax.experimental.mosaic.gpu.layout_inference.extract_assignment_candidates_from_reduce_equation", "jax.experimental.mosaic.gpu.layout_inference.find_assignments_for", "jax.experimental.mosaic.gpu.layout_inference.infer_layout", "jax.experimental.mosaic.gpu.layout_inference.is_terminator", "jax.experimental.mosaic.gpu.layout_inference.is_valid_register_layout_assignment", "jax.experimental.mosaic.gpu.layout_inference.is_valid_smem_layout_assignment", "jax.experimental.mosaic.gpu.layout_inference.is_valid_tmem_layout_assignment", "jax.experimental.mosaic.gpu.layout_inference.is_vector", "jax.experimental.mosaic.gpu.layout_inference.prime_decomposition", "jax.experimental.mosaic.gpu.layout_inference.producer_result", "jax.experimental.mosaic.gpu.layout_inference.traverse_op", "jax.experimental.mosaic.gpu.layout_inference.vector_value_sites" ], "n_typable": 58, "n_typed": 54, "n_any": 0, "n_untyped": 4, "n_functions": 18, "n_function_overloads": 18, "n_function_params": 30, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 2, "n_classes": 4, "n_attrs": 15, "n_properties": 3, "n_type_ignores": 0 }, { "path": "jax/experimental/mosaic/gpu/constraints.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.mosaic.gpu.constraints.non_splat_variables", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.constraints.saturate_divides_constraints_for_equal_vars", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.constraints.Equals", "methods": [ { "kind": "function", "name": "Equals.holds", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Equals.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Equals.lhs", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Equals.rhs", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 2, "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.constraints.NotOfType", "methods": [ { "kind": "function", "name": "NotOfType.holds", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "NotOfType.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "NotOfType.expr", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "NotOfType.type", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 2, "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.constraints.IsValidMmaTiling", "methods": [ { "kind": "function", "name": "IsValidMmaTiling.holds", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "IsValidMmaTiling.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "IsValidMmaTiling.expr", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "IsValidMmaTiling.bitwidth", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 2, "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.constraints.reduce_broadcast_expression", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.constraints.Constraint", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.constraints.ConstraintSystem", "methods": [ { "kind": "function", "name": "ConstraintSystem.unknowns", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ConstraintSystem.__and__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ConstraintSystem.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "ConstraintSystem.assignments", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ConstraintSystem.constraints", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 4, "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 1, "n_attrs": 2, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.constraints.Expression", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.constraints.SMEMTiling", "methods": [ { "kind": "function", "name": "SMEMTiling.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "SMEMTiling.value", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 1, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.constraints.BroadcastInDim", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "BroadcastInDim.expression", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BroadcastInDim.axes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BroadcastInDim.shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.constraints.reduce_reshape_expression", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.constraints.merge_divides_constraints", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.constraints.Reduce", "methods": [ { "kind": "function", "name": "Reduce.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Reduce.expression", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Reduce.axes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.constraints.saturate_distinct_from_splat", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.constraints.reduce", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.constraints.Constant", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.constraints.Reshape", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "Reshape.expression", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Reshape.source_shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Reshape.target_shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.constraints.Transpose", "methods": [ { "kind": "function", "name": "Transpose.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Transpose.expression", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 1, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.constraints.reduce_expression", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.constraints.IsTransferable", "methods": [ { "kind": "function", "name": "IsTransferable.supported_tmem_transfers", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "IsTransferable._is_valid_tmem_transfer", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "IsTransferable._is_valid_smem_transfer", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "IsTransferable.holds", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "IsTransferable.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "IsTransferable.source", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "IsTransferable.target", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "IsTransferable.shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 10, "n_typed": 9, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 5, "n_attrs": 3, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.constraints.RegisterLayout", "methods": [ { "kind": "function", "name": "RegisterLayout.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "RegisterLayout.value", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 1, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.constraints.Relayout", "methods": [ { "kind": "function", "name": "Relayout.holds", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Relayout.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Relayout.source", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Relayout.target", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Relayout.bitwidth", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 2, "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.constraints.reduce_constraint", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.constraints.compute_transitively_equal_vars", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.constraints.TMEMLayout", "methods": [ { "kind": "function", "name": "TMEMLayout.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "TMEMLayout.value", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 1, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.constraints.Variable", "methods": [ { "kind": "function", "name": "Variable.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Variable.key", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 1, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.constraints.Unsatisfiable", "methods": [ { "kind": "function", "name": "Unsatisfiable.__and__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [], "n_typable": 2, "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 1, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.constraints.Divides", "methods": [ { "kind": "function", "name": "Divides.holds", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Divides.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Divides.expr", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Divides.tiling_multiple", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 2, "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.constraints.reduce_transpose_expression", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [], "name": "jax.experimental.mosaic.gpu.constraints", "names": [ "jax.experimental.mosaic.gpu.constraints.BroadcastInDim", "jax.experimental.mosaic.gpu.constraints.Constant", "jax.experimental.mosaic.gpu.constraints.Constraint", "jax.experimental.mosaic.gpu.constraints.ConstraintSystem", "jax.experimental.mosaic.gpu.constraints.Divides", "jax.experimental.mosaic.gpu.constraints.Equals", "jax.experimental.mosaic.gpu.constraints.Expression", "jax.experimental.mosaic.gpu.constraints.IsTransferable", "jax.experimental.mosaic.gpu.constraints.IsValidMmaTiling", "jax.experimental.mosaic.gpu.constraints.NotOfType", "jax.experimental.mosaic.gpu.constraints.Reduce", "jax.experimental.mosaic.gpu.constraints.RegisterLayout", "jax.experimental.mosaic.gpu.constraints.Relayout", "jax.experimental.mosaic.gpu.constraints.Reshape", "jax.experimental.mosaic.gpu.constraints.SMEMTiling", "jax.experimental.mosaic.gpu.constraints.TMEMLayout", "jax.experimental.mosaic.gpu.constraints.Transpose", "jax.experimental.mosaic.gpu.constraints.Unsatisfiable", "jax.experimental.mosaic.gpu.constraints.Variable", "jax.experimental.mosaic.gpu.constraints.compute_transitively_equal_vars", "jax.experimental.mosaic.gpu.constraints.merge_divides_constraints", "jax.experimental.mosaic.gpu.constraints.non_splat_variables", "jax.experimental.mosaic.gpu.constraints.reduce", "jax.experimental.mosaic.gpu.constraints.reduce_broadcast_expression", "jax.experimental.mosaic.gpu.constraints.reduce_constraint", "jax.experimental.mosaic.gpu.constraints.reduce_expression", "jax.experimental.mosaic.gpu.constraints.reduce_reshape_expression", "jax.experimental.mosaic.gpu.constraints.reduce_transpose_expression", "jax.experimental.mosaic.gpu.constraints.saturate_distinct_from_splat", "jax.experimental.mosaic.gpu.constraints.saturate_divides_constraints_for_equal_vars" ], "n_typable": 59, "n_typed": 46, "n_any": 0, "n_untyped": 13, "n_functions": 11, "n_function_overloads": 11, "n_function_params": 16, "n_methods": 25, "n_method_overloads": 25, "n_method_params": 7, "n_classes": 17, "n_attrs": 31, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/mosaic/gpu/inference_utils.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.in_tmem_layouts", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.attr_element", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.should_have_out_tmem_layout", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.inference_utils.in_transforms_for_operand", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.is_mma_layout", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.should_have_in_tmem_layout", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.out_transforms", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.inference_utils.MlirOperation", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.has_out_layouts_set", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.out_tmem_layouts", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.has_any_layout_set", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.should_have_transforms", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.inference_utils.in_layout_for_operand", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.in_layouts", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.is_transformable_smem_memref", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.in_transforms", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.has_in_transforms_set", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.out_layouts", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.should_have_in_transforms", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.has_out_tmem_layouts_set", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.has_in_tmem_layouts_set", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.should_have_in_layout", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.should_have_out_transforms", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.should_have_tmem_layout", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.should_have_out_layout", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.has_in_layouts_set", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.should_have_layout", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.inference_utils.has_out_transforms_set", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.mosaic.gpu.inference_utils", "names": [ "jax.experimental.mosaic.gpu.inference_utils.MlirOperation", "jax.experimental.mosaic.gpu.inference_utils.attr_element", "jax.experimental.mosaic.gpu.inference_utils.has_any_layout_set", "jax.experimental.mosaic.gpu.inference_utils.has_in_layouts_set", "jax.experimental.mosaic.gpu.inference_utils.has_in_tmem_layouts_set", "jax.experimental.mosaic.gpu.inference_utils.has_in_transforms_set", "jax.experimental.mosaic.gpu.inference_utils.has_out_layouts_set", "jax.experimental.mosaic.gpu.inference_utils.has_out_tmem_layouts_set", "jax.experimental.mosaic.gpu.inference_utils.has_out_transforms_set", "jax.experimental.mosaic.gpu.inference_utils.in_layout_for_operand", "jax.experimental.mosaic.gpu.inference_utils.in_layouts", "jax.experimental.mosaic.gpu.inference_utils.in_tmem_layouts", "jax.experimental.mosaic.gpu.inference_utils.in_transforms", "jax.experimental.mosaic.gpu.inference_utils.in_transforms_for_operand", "jax.experimental.mosaic.gpu.inference_utils.is_mma_layout", "jax.experimental.mosaic.gpu.inference_utils.is_transformable_smem_memref", "jax.experimental.mosaic.gpu.inference_utils.out_layouts", "jax.experimental.mosaic.gpu.inference_utils.out_tmem_layouts", "jax.experimental.mosaic.gpu.inference_utils.out_transforms", "jax.experimental.mosaic.gpu.inference_utils.should_have_in_layout", "jax.experimental.mosaic.gpu.inference_utils.should_have_in_tmem_layout", "jax.experimental.mosaic.gpu.inference_utils.should_have_in_transforms", "jax.experimental.mosaic.gpu.inference_utils.should_have_layout", "jax.experimental.mosaic.gpu.inference_utils.should_have_out_layout", "jax.experimental.mosaic.gpu.inference_utils.should_have_out_tmem_layout", "jax.experimental.mosaic.gpu.inference_utils.should_have_out_transforms", "jax.experimental.mosaic.gpu.inference_utils.should_have_tmem_layout", "jax.experimental.mosaic.gpu.inference_utils.should_have_transforms" ], "n_typable": 55, "n_typed": 53, "n_any": 0, "n_untyped": 2, "n_functions": 25, "n_function_overloads": 25, "n_function_params": 27, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 3, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/mosaic/gpu/mma_utils.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.mosaic.gpu.mma_utils.create_descriptor", "n_typed": 7, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.mma_utils.Dim", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "Dim.K", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Dim.MN", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.mma_utils.tiled_memref_shape", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.mma_utils.encode_descriptor", "n_typed": 5, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.mma_utils.encode_addr", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.mosaic.gpu.mma_utils", "names": [ "jax.experimental.mosaic.gpu.mma_utils.Dim", "jax.experimental.mosaic.gpu.mma_utils.create_descriptor", "jax.experimental.mosaic.gpu.mma_utils.encode_addr", "jax.experimental.mosaic.gpu.mma_utils.encode_descriptor", "jax.experimental.mosaic.gpu.mma_utils.tiled_memref_shape" ], "n_typable": 19, "n_typed": 14, "n_any": 0, "n_untyped": 5, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 15, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 2, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/mosaic/gpu/profiler.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.mosaic.gpu.profiler.P", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.profiler.Cupti", "methods": [ { "kind": "function", "name": "Cupti.measure", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Cupti.finalize", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 4, "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 3, "n_attrs": 1, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.profiler.T", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.profiler.ProfilerSpec", "methods": [ { "kind": "function", "name": "ProfilerSpec.__init__", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "ProfilerSpec._num_warpgroups", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "ProfilerSpec.mlir_buffer_type", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "ProfilerSpec.jax_buffer_type", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "ProfilerSpec.smem_i32_elements", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ProfilerSpec.smem_bytes", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ProfilerSpec.intern_name", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ProfilerSpec.dump", "n_typed": 2, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "ProfilerSpec.ENTER", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ProfilerSpec.EXIT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ProfilerSpec.entries_per_warpgroup", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "ProfilerSpec.interned_names", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "ProfilerSpec.dump_path", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 26, "n_typed": 18, "n_any": 0, "n_untyped": 8, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 8, "n_method_overloads": 8, "n_method_params": 15, "n_attrs": 5, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.profiler.measure", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 4, "n_typable": 4, "n_params": 3, "n_function_overloads": 4, "n_function_params": 3 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.profiler.OnDeviceProfiler", "methods": [ { "kind": "function", "name": "OnDeviceProfiler.__init__", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "OnDeviceProfiler._profiler_ctx", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "OnDeviceProfiler.record", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "OnDeviceProfiler.finalize", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "OnDeviceProfiler.spec", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "OnDeviceProfiler.entries_per_wg", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "OnDeviceProfiler.wrap_in_custom_primitive", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "OnDeviceProfiler.ctx", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 15, "n_typed": 7, "n_any": 0, "n_untyped": 8, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 7, "n_attrs": 4, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.profiler.mosaic_gpu_lib", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [ { "kind": "type", "rules": [ "assignment" ] } ], "name": "jax.experimental.mosaic.gpu.profiler", "names": [ "jax.experimental.mosaic.gpu.profiler.Cupti", "jax.experimental.mosaic.gpu.profiler.OnDeviceProfiler", "jax.experimental.mosaic.gpu.profiler.P", "jax.experimental.mosaic.gpu.profiler.ProfilerSpec", "jax.experimental.mosaic.gpu.profiler.T", "jax.experimental.mosaic.gpu.profiler.measure", "jax.experimental.mosaic.gpu.profiler.mosaic_gpu_lib" ], "n_typable": 49, "n_typed": 31, "n_any": 0, "n_untyped": 18, "n_functions": 1, "n_function_overloads": 4, "n_function_params": 3, "n_methods": 13, "n_method_overloads": 13, "n_method_params": 25, "n_classes": 3, "n_attrs": 13, "n_properties": 0, "n_type_ignores": 1 }, { "path": "jax/experimental/mosaic/gpu/tcgen05.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.mosaic.gpu.tcgen05.TCGEN05_SMEM_DESCRIPTOR_BIT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.tcgen05.TMEMRef", "methods": [ { "kind": "function", "name": "TMEMRef.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "TMEMRef.from_alloc", "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "TMEMRef.slice", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TMEMRef.load", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "TMEMRef.store", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TMEMRef._debug_print", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "TMEMRef.packing", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "TMEMRef.address", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TMEMRef.shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TMEMRef.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TMEMRef.layout", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 16, "n_typed": 12, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 9, "n_attrs": 4, "n_properties": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.tmem_relinquish_alloc_permit", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.commit_arrive", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.mma", "n_typed": 11, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 11, "n_params": 10, "n_function_overloads": 1, "n_function_params": 10 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.async_copy_scales_smem_to_tmem", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.tcgen05.TMEM_ROWS", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.async_copy_sparse_metadata_smem_to_tmem", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.create_instr_descriptor", "n_typed": 6, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.commit_tmem", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.tmem_dealloc", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.tmem_half_lane_layout", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.scales_layout", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.tmem_m64_collective_layout", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.create_scaled_f8f6f4_instr_descriptor", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.tmem_alloc", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.wait_load_tmem", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.async_copy_smem_to_tmem", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.sparse_meta_layout", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "attr", "name": "jax.experimental.mosaic.gpu.tcgen05.TMEM_MAX_COLS", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "class", "name": "jax.experimental.mosaic.gpu.tcgen05.TMEMLayout", "methods": [ { "kind": "function", "name": "TMEMLayout.check_type", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "TMEMLayout.cols_in_shape", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "TMEMLayout.canonicalize", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "TMEMLayout.as_tiled_layout", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [], "n_typable": 8, "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 4, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.fa_m64_collective_layout", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.tmem_alloc_exact_ncols", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.tmem_default_layout", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.mosaic.gpu.tcgen05.create_scaled_f4_instr_descriptor", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [], "name": "jax.experimental.mosaic.gpu.tcgen05", "names": [ "jax.experimental.mosaic.gpu.tcgen05.TCGEN05_SMEM_DESCRIPTOR_BIT", "jax.experimental.mosaic.gpu.tcgen05.TMEMLayout", "jax.experimental.mosaic.gpu.tcgen05.TMEMRef", "jax.experimental.mosaic.gpu.tcgen05.TMEM_MAX_COLS", "jax.experimental.mosaic.gpu.tcgen05.TMEM_ROWS", "jax.experimental.mosaic.gpu.tcgen05.async_copy_scales_smem_to_tmem", "jax.experimental.mosaic.gpu.tcgen05.async_copy_smem_to_tmem", "jax.experimental.mosaic.gpu.tcgen05.async_copy_sparse_metadata_smem_to_tmem", "jax.experimental.mosaic.gpu.tcgen05.commit_arrive", "jax.experimental.mosaic.gpu.tcgen05.commit_tmem", "jax.experimental.mosaic.gpu.tcgen05.create_instr_descriptor", "jax.experimental.mosaic.gpu.tcgen05.create_scaled_f4_instr_descriptor", "jax.experimental.mosaic.gpu.tcgen05.create_scaled_f8f6f4_instr_descriptor", "jax.experimental.mosaic.gpu.tcgen05.fa_m64_collective_layout", "jax.experimental.mosaic.gpu.tcgen05.mma", "jax.experimental.mosaic.gpu.tcgen05.scales_layout", "jax.experimental.mosaic.gpu.tcgen05.sparse_meta_layout", "jax.experimental.mosaic.gpu.tcgen05.tmem_alloc", "jax.experimental.mosaic.gpu.tcgen05.tmem_alloc_exact_ncols", "jax.experimental.mosaic.gpu.tcgen05.tmem_dealloc", "jax.experimental.mosaic.gpu.tcgen05.tmem_default_layout", "jax.experimental.mosaic.gpu.tcgen05.tmem_half_lane_layout", "jax.experimental.mosaic.gpu.tcgen05.tmem_m64_collective_layout", "jax.experimental.mosaic.gpu.tcgen05.tmem_relinquish_alloc_permit", "jax.experimental.mosaic.gpu.tcgen05.wait_load_tmem" ], "n_typable": 95, "n_typed": 84, "n_any": 0, "n_untyped": 11, "n_functions": 20, "n_function_overloads": 20, "n_function_params": 51, "n_methods": 10, "n_method_overloads": 10, "n_method_params": 13, "n_classes": 2, "n_attrs": 7, "n_properties": 1, "n_type_ignores": 0 }, { "path": "jax/experimental/multihost_utils.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.multihost_utils.host_local_array_to_global_array_impl", "n_typed": 1, "n_any": 2, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.multihost_utils.gtl_abstract_eval", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.multihost_utils.global_array_to_host_local_array", "n_typed": 1, "n_any": 2, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.multihost_utils.process_allgather", "n_typed": 1, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.multihost_utils.reached_preemption_sync_point", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.multihost_utils.global_array_to_host_local_array_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.multihost_utils.global_array_to_host_local_array_impl", "n_typed": 1, "n_any": 2, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.multihost_utils.ltg_batcher", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.experimental.multihost_utils.ltg_abstract_eval", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.multihost_utils.broadcast_one_to_all", "n_typed": 1, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.multihost_utils.assert_equal", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.experimental.multihost_utils.live_devices", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.multihost_utils.host_local_array_to_global_array_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.multihost_utils.host_local_array_to_global_array", "n_typed": 1, "n_any": 2, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.multihost_utils.sync_global_devices", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.multihost_utils", "names": [ "jax.experimental.multihost_utils.assert_equal", "jax.experimental.multihost_utils.broadcast_one_to_all", "jax.experimental.multihost_utils.global_array_to_host_local_array", "jax.experimental.multihost_utils.global_array_to_host_local_array_impl", "jax.experimental.multihost_utils.global_array_to_host_local_array_p", "jax.experimental.multihost_utils.gtl_abstract_eval", "jax.experimental.multihost_utils.host_local_array_to_global_array", "jax.experimental.multihost_utils.host_local_array_to_global_array_impl", "jax.experimental.multihost_utils.host_local_array_to_global_array_p", "jax.experimental.multihost_utils.live_devices", "jax.experimental.multihost_utils.ltg_abstract_eval", "jax.experimental.multihost_utils.ltg_batcher", "jax.experimental.multihost_utils.process_allgather", "jax.experimental.multihost_utils.reached_preemption_sync_point", "jax.experimental.multihost_utils.sync_global_devices" ], "n_typable": 47, "n_typed": 10, "n_any": 12, "n_untyped": 25, "n_functions": 12, "n_function_overloads": 12, "n_function_params": 32, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 3, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/ode.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.ode.abs2", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.ode.initial_step_size", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax.experimental.ode.mean_error_ratio", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.experimental.ode.optimal_step_size", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.experimental.ode.ravel_first_arg", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.ode.fit_4th_order_polynomial", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.experimental.ode.runge_kutta_step", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.experimental.ode.ravel_first_arg_", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.ode.odeint", "n_typed": 0, "n_any": 0, "n_untyped": 9, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.experimental.ode.interp_fit_dopri", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax.experimental.ode", "names": [ "jax.experimental.ode.abs2", "jax.experimental.ode.fit_4th_order_polynomial", "jax.experimental.ode.initial_step_size", "jax.experimental.ode.interp_fit_dopri", "jax.experimental.ode.mean_error_ratio", "jax.experimental.ode.odeint", "jax.experimental.ode.optimal_step_size", "jax.experimental.ode.ravel_first_arg", "jax.experimental.ode.ravel_first_arg_", "jax.experimental.ode.runge_kutta_step" ], "n_typable": 59, "n_typed": 2, "n_any": 0, "n_untyped": 57, "n_functions": 10, "n_function_overloads": 10, "n_function_params": 49, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/helpers.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.helpers.loop", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.helpers.kernel", "n_typed": 10, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 11, "n_params": 10, "n_function_overloads": 1, "n_function_params": 10 }, { "kind": "function", "name": "jax._src.pallas.helpers.empty_ref_like", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.pallas.helpers.empty_like", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.pallas.helpers.debug_checks_enabled", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.pallas.helpers.when", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax._src.pallas.helpers.empty", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.pallas.helpers.with_scoped", "n_typed": 1, "n_any": 2, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.helpers.debug_check", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [], "name": "jax._src.pallas.helpers", "names": [ "jax._src.pallas.helpers.debug_check", "jax._src.pallas.helpers.debug_checks_enabled", "jax._src.pallas.helpers.empty", "jax._src.pallas.helpers.empty_like", "jax._src.pallas.helpers.empty_ref_like", "jax._src.pallas.helpers.kernel", "jax._src.pallas.helpers.loop", "jax._src.pallas.helpers.when", "jax._src.pallas.helpers.with_scoped" ], "n_typable": 31, "n_typed": 22, "n_any": 2, "n_untyped": 7, "n_functions": 8, "n_function_overloads": 8, "n_function_params": 22, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/core.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.pallas.core.Squeezed", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.core.BoundedSlice", "methods": [ { "kind": "function", "name": "BoundedSlice.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "BoundedSlice.block_size", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 1, "n_properties": 0 }, { "kind": "function", "name": "jax._src.pallas.core.lower_as_mlir", "n_typed": 1, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "class", "name": "jax._src.pallas.core.MemorySpace", "methods": [ { "kind": "function", "name": "MemorySpace.from_type", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "MemorySpace.__call__", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "MemorySpace.__str__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "MemorySpace.ANY", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MemorySpace.ERROR", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MemorySpace.INDEX", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MemorySpace.KEY", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MemorySpace.HOST", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 6, "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 3, "n_attrs": 5, "n_properties": 0 }, { "kind": "function", "name": "jax._src.pallas.core.core_map", "n_typed": 7, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "attr", "name": "jax._src.pallas.core.no_block_spec", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "class", "name": "jax._src.pallas.core.Buffered", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "Buffered.buffer_count", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Buffered.use_lookahead", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "attr", "name": "jax._src.pallas.core.BlockDim", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.pallas.core.squeezed", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "class", "name": "jax._src.pallas.core.CostEstimate", "methods": [ { "kind": "function", "name": "CostEstimate.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "CostEstimate.to_json", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "CostEstimate.flops", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CostEstimate.transcendentals", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CostEstimate.bytes_accessed", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CostEstimate.remote_bytes_transferred", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 2, "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 0, "n_attrs": 4, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.core.CompilerParams", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.core.MemoryRef", "methods": [ { "kind": "function", "name": "MemoryRef.get_array_aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "MemoryRef.get_ref_aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "MemoryRef.__lt__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "MemoryRef.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "MemoryRef.shape", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "MemoryRef.inner_aval", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MemoryRef.memory_space", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 6, "n_typed": 2, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 1, "n_attrs": 2, "n_properties": 2 }, { "kind": "class", "name": "jax._src.pallas.core.GridSpec", "methods": [ { "kind": "function", "name": "GridSpec.__init__", "n_typed": 2, "n_any": 2, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "GridSpec._make_scalar_ref_aval", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "GridSpec.grid", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GridSpec.grid_names", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GridSpec.in_specs", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GridSpec.out_specs", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GridSpec.scratch_shapes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 7, "n_typed": 2, "n_any": 2, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 5, "n_attrs": 5, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.core.Blocked", "methods": [ { "kind": "function", "name": "Blocked.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Blocked.block_size", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 1, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.core.semaphore", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.core.BlockSpec", "methods": [ { "kind": "function", "name": "BlockSpec.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BlockSpec.to_block_mapping", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "BlockSpec.block_shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSpec.index_map", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSpec.memory_space", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSpec.pipeline_mode", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSpec.replace", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 9, "n_typed": 8, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 7, "n_attrs": 5, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.core.Element", "methods": [ { "kind": "function", "name": "Element.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Element.block_size", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Element.padding", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 } ], "type_ignores": [ { "kind": "type", "rules": [ "method-assign" ] } ], "name": "jax._src.pallas.core", "names": [ "jax._src.pallas.core.BlockDim", "jax._src.pallas.core.BlockSpec", "jax._src.pallas.core.Blocked", "jax._src.pallas.core.BoundedSlice", "jax._src.pallas.core.Buffered", "jax._src.pallas.core.CompilerParams", "jax._src.pallas.core.CostEstimate", "jax._src.pallas.core.Element", "jax._src.pallas.core.GridSpec", "jax._src.pallas.core.MemoryRef", "jax._src.pallas.core.MemorySpace", "jax._src.pallas.core.Squeezed", "jax._src.pallas.core.core_map", "jax._src.pallas.core.lower_as_mlir", "jax._src.pallas.core.no_block_spec", "jax._src.pallas.core.semaphore", "jax._src.pallas.core.squeezed" ], "n_typable": 53, "n_typed": 27, "n_any": 2, "n_untyped": 24, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 15, "n_methods": 15, "n_method_overloads": 15, "n_method_params": 16, "n_classes": 12, "n_attrs": 30, "n_properties": 2, "n_type_ignores": 1 }, { "path": "jax/_src/pallas/cost_estimate.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.cost_estimate.estimate_cost", "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.pallas.cost_estimate", "names": [ "jax._src.pallas.cost_estimate.estimate_cost" ], "n_typable": 4, "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 3, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/primitives.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.pallas.primitives.DeviceIdType", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "DeviceIdType.MESH", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DeviceIdType.LOGICAL", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax._src.pallas.primitives.dot", "n_typed": 3, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.pallas.primitives.semaphore_signal", "n_typed": 4, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.pallas.primitives.semaphore_read", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.pallas.primitives.debug_print", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.primitives.multiple_of", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.primitives.num_programs", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.pallas.primitives.reciprocal", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.primitives.delay", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.pallas.primitives.program_id", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.pallas.primitives.semaphore_wait", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.primitives.get_global", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.pallas.primitives.run_scoped", "n_typed": 2, "n_any": 3, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.primitives.atomic_or", "n_typed": 1, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.primitives.atomic_min", "n_typed": 1, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.primitives.max_contiguous", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.primitives.atomic_xor", "n_typed": 1, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.primitives.atomic_cas", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.primitives.atomic_max", "n_typed": 1, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.primitives.atomic_add", "n_typed": 1, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.primitives.atomic_and", "n_typed": 1, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.primitives.atomic_xchg", "n_typed": 1, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [ { "kind": "type", "rules": [ "method-assign" ] } ], "name": "jax._src.pallas.primitives", "names": [ "jax._src.pallas.primitives.DeviceIdType", "jax._src.pallas.primitives.atomic_add", "jax._src.pallas.primitives.atomic_and", "jax._src.pallas.primitives.atomic_cas", "jax._src.pallas.primitives.atomic_max", "jax._src.pallas.primitives.atomic_min", "jax._src.pallas.primitives.atomic_or", "jax._src.pallas.primitives.atomic_xchg", "jax._src.pallas.primitives.atomic_xor", "jax._src.pallas.primitives.debug_print", "jax._src.pallas.primitives.delay", "jax._src.pallas.primitives.dot", "jax._src.pallas.primitives.get_global", "jax._src.pallas.primitives.max_contiguous", "jax._src.pallas.primitives.multiple_of", "jax._src.pallas.primitives.num_programs", "jax._src.pallas.primitives.program_id", "jax._src.pallas.primitives.reciprocal", "jax._src.pallas.primitives.run_scoped", "jax._src.pallas.primitives.semaphore_read", "jax._src.pallas.primitives.semaphore_signal", "jax._src.pallas.primitives.semaphore_wait" ], "n_typable": 83, "n_typed": 32, "n_any": 3, "n_untyped": 48, "n_functions": 21, "n_function_overloads": 21, "n_function_params": 62, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 2, "n_properties": 0, "n_type_ignores": 1 }, { "path": "jax/_src/pallas/utils.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.utils.strides_from_shape", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.pallas.utils.cdiv", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 4, "n_typable": 3, "n_params": 2, "n_function_overloads": 4, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.utils.next_power_of_2", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax._src.pallas.utils", "names": [ "jax._src.pallas.utils.cdiv", "jax._src.pallas.utils.next_power_of_2", "jax._src.pallas.utils.strides_from_shape" ], "n_typable": 7, "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_functions": 3, "n_function_overloads": 6, "n_function_params": 4, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/__init__.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.pallas.enable_debug_checks", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.pallas.ds", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [], "name": "jax.experimental.pallas", "names": [ "jax.experimental.pallas.ds", "jax.experimental.pallas.enable_debug_checks" ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 2, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/indexing.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.indexing.dslice", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax._src.indexing.Slice", "methods": [ { "kind": "function", "name": "Slice.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Slice.tree_flatten", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Slice.tree_unflatten", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Slice.from_slice", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [ { "kind": "property", "name": "Slice.is_dynamic_start", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Slice.is_dynamic_size", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "Slice.start", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Slice.size", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Slice.stride", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 10, "n_typed": 4, "n_any": 0, "n_untyped": 6, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 4, "n_attrs": 3, "n_properties": 2 } ], "type_ignores": [], "name": "jax._src.indexing", "names": [ "jax._src.indexing.Slice", "jax._src.indexing.dslice" ], "n_typable": 14, "n_typed": 8, "n_any": 0, "n_untyped": 6, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 3, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 4, "n_classes": 1, "n_attrs": 3, "n_properties": 2, "n_type_ignores": 0 }, { "path": "jax/_src/state/discharge.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.state.discharge.run_state", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [ { "kind": "type", "rules": null }, { "kind": "pyrefly", "rules": [ "bad-assignment" ] } ], "name": "jax._src.state.discharge", "names": [ "jax._src.state.discharge.run_state" ], "n_typable": 2, "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 1, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 2 }, { "path": "jax/_src/pallas/pallas_call.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.pallas_call.pallas_call", "n_typed": 13, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 15, "n_params": 14, "n_function_overloads": 1, "n_function_params": 14 }, { "kind": "attr", "name": "jax._src.pallas.pallas_call.pallas_call_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [ { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null } ], "name": "jax._src.pallas.pallas_call", "names": [ "jax._src.pallas.pallas_call.pallas_call", "jax._src.pallas.pallas_call.pallas_call_p" ], "n_typable": 16, "n_typed": 13, "n_any": 2, "n_untyped": 1, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 14, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 7 }, { "path": "jax/_src/state/primitives.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.state.primitives.broadcast_to", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.state.primitives.ref_addupdate", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.state.primitives.ref_swap", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.state.primitives.ref_set", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.state.primitives.ref_get", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [ { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null } ], "name": "jax._src.state.primitives", "names": [ "jax._src.state.primitives.broadcast_to", "jax._src.state.primitives.ref_addupdate", "jax._src.state.primitives.ref_get", "jax._src.state.primitives.ref_set", "jax._src.state.primitives.ref_swap" ], "n_typable": 19, "n_typed": 19, "n_any": 0, "n_untyped": 0, "n_functions": 5, "n_function_overloads": 5, "n_function_params": 14, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 7 }, { "path": "jax/_src/pallas/fuser/block_spec.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.fuser.block_spec.get_fusion_values", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.fuser.block_spec.pull_block_spec", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.fuser.block_spec.push_block_spec", "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.fuser.block_spec.make_scalar_prefetch_handler", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax._src.pallas.fuser.block_spec", "names": [ "jax._src.pallas.fuser.block_spec.get_fusion_values", "jax._src.pallas.fuser.block_spec.make_scalar_prefetch_handler", "jax._src.pallas.fuser.block_spec.pull_block_spec", "jax._src.pallas.fuser.block_spec.push_block_spec" ], "n_typable": 15, "n_typed": 7, "n_any": 0, "n_untyped": 8, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 11, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/fuser/custom_fusion_lib.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.pallas.fuser.custom_fusion_lib.custom_fusion", "methods": [ { "kind": "function", "name": "custom_fusion.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "custom_fusion.def_pallas_impl", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "custom_fusion.def_pull_block_spec", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "custom_fusion.def_push_block_spec", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "custom_fusion.def_eval_rule", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "custom_fusion.__call__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "custom_fusion.fun", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "custom_fusion.eval_rule", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "custom_fusion.pull_block_spec_rule", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "custom_fusion.push_block_spec_rule", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "custom_fusion.pallas_impl", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 18, "n_typed": 9, "n_any": 0, "n_untyped": 9, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 7, "n_attrs": 5, "n_properties": 0 } ], "type_ignores": [ { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": [ "arg-type" ] } ], "name": "jax._src.pallas.fuser.custom_fusion_lib", "names": [ "jax._src.pallas.fuser.custom_fusion_lib.custom_fusion" ], "n_typable": 18, "n_typed": 9, "n_any": 0, "n_untyped": 9, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 7, "n_classes": 1, "n_attrs": 5, "n_properties": 0, "n_type_ignores": 3 }, { "path": "jax/_src/pallas/fuser/custom_evaluate.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.fuser.custom_evaluate.evaluate", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [], "name": "jax._src.pallas.fuser.custom_evaluate", "names": [ "jax._src.pallas.fuser.custom_evaluate.evaluate" ], "n_typable": 3, "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 2, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/fuser/fusible.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.fuser.fusible.fusible", "n_typed": 0, "n_any": 1, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [ { "kind": "type", "rules": null } ], "name": "jax._src.pallas.fuser.fusible", "names": [ "jax._src.pallas.fuser.fusible.fusible" ], "n_typable": 3, "n_typed": 0, "n_any": 1, "n_untyped": 2, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 2, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 1 }, { "path": "jax/_src/pallas/fuser/fusion.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.pallas.fuser.fusion.Fusion", "methods": [ { "kind": "function", "name": "Fusion.__call__", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [ { "kind": "property", "name": "Fusion.shape", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Fusion.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Fusion.type", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Fusion.in_shape", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Fusion.in_dtype", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "Fusion.func", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Fusion.in_type", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Fusion.out_type", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 8, "n_typed": 3, "n_any": 0, "n_untyped": 5, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 2, "n_attrs": 3, "n_properties": 5 } ], "type_ignores": [], "name": "jax._src.pallas.fuser.fusion", "names": [ "jax._src.pallas.fuser.fusion.Fusion" ], "n_typable": 8, "n_typed": 3, "n_any": 0, "n_untyped": 5, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 2, "n_classes": 1, "n_attrs": 3, "n_properties": 5, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/fuser/jaxpr_fusion.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.fuser.jaxpr_fusion.fuse", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.pallas.fuser.jaxpr_fusion", "names": [ "jax._src.pallas.fuser.jaxpr_fusion.fuse" ], "n_typable": 4, "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 3, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/mosaic_gpu/core.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.core.transpose_ref", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.core.kernel", "n_typed": 10, "n_any": 2, "n_untyped": 1, "n_overloads": 1, "n_typable": 13, "n_params": 12, "n_function_overloads": 1, "n_function_params": 12 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.core.PeerMemRef", "methods": [ { "kind": "function", "name": "PeerMemRef.undo", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "PeerMemRef.transform_type", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "PeerMemRef.commute_ndindexer", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "PeerMemRef.device_id", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "PeerMemRef.device_id_type", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 7, "n_typed": 5, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 4, "n_attrs": 2, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.core.RefUnion", "methods": [ { "kind": "function", "name": "RefUnion.__init__", "n_typed": 0, "n_any": 1, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "RefUnion.get_ref_aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "RefUnion.refs", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 3, "n_typed": 1, "n_any": 1, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 1, "n_attrs": 1, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.core.SwizzleTransform", "methods": [ { "kind": "function", "name": "SwizzleTransform.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "SwizzleTransform.transform_type", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SwizzleTransform.undo", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "SwizzleTransform.swizzle", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 5, "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 2, "n_attrs": 1, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.core.Barrier", "methods": [ { "kind": "function", "name": "Barrier.get_array_aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Barrier.get_ref_aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Barrier.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Barrier.num_arrivals", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Barrier.num_barriers", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Barrier.orders_tensor_core", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 3, "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.core.transform_ref", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.core.multicast_ref", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.core.ClusterBarrier", "methods": [ { "kind": "function", "name": "ClusterBarrier.get_array_aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ClusterBarrier.get_ref_aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "ClusterBarrier.collective_axes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ClusterBarrier.num_barriers", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ClusterBarrier.num_arrivals", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ClusterBarrier.orders_tensor_core", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 2, "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 0, "n_attrs": 4, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.core.MemorySpace", "methods": [ { "kind": "function", "name": "MemorySpace.__str__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "MemorySpace.__call__", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "MemorySpace.GMEM", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MemorySpace.SMEM", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MemorySpace.TMEM", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MemorySpace.REGS", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 8, "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 6, "n_attrs": 4, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.core.WGMMAAccumulatorRef", "methods": [ { "kind": "function", "name": "WGMMAAccumulatorRef.get_ref_aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "WGMMAAccumulatorRef.init", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "WGMMAAccumulatorRef.shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "WGMMAAccumulatorRef.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "WGMMAAccumulatorRef._init", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 3, "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 1, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.core.unswizzle_ref", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.core.Layout", "methods": [ { "kind": "function", "name": "Layout.__call__", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Layout.to_mgpu", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Layout.WGMMA", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Layout.WGMMA_8BIT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Layout.WGMMA_UPCAST_2X", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Layout.WGMMA_UPCAST_4X", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Layout.WGMMA_TRANSPOSED", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Layout.WG_SPLAT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Layout.WG_STRIDED", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Layout.TILED", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Layout.TCGEN05", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Layout.TCGEN05_TRANSPOSED", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Layout.TCGEN05_M64_COLLECTIVE", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Layout.TCGEN05_TMEM_NATIVE", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Layout.TCGEN05_M64_COLLECTIVE_NATIVE", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Layout.SMEM_GMEM_COPY", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Layout.TMA_GATHER_INDICES", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Layout._WGMMA_ACC_32BIT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 6, "n_typed": 2, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 4, "n_attrs": 16, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.core.TilingTransform", "methods": [ { "kind": "function", "name": "TilingTransform.transform_type", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TilingTransform.undo", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "TilingTransform.tiling", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 4, "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 2, "n_attrs": 1, "n_properties": 0 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.core.layout_cast", "n_typed": 1, "n_any": 1, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.core.Mesh", "methods": [ { "kind": "function", "name": "Mesh.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Mesh.discharges_effect", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "Mesh.default_memory_space", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Mesh.shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "Mesh.grid", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Mesh.grid_names", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Mesh.cluster", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Mesh.cluster_names", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Mesh.num_threads", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Mesh.thread_name", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Mesh.kernel_name", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 5, "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 1, "n_attrs": 7, "n_properties": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.core.remote_ref", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.core.TryClusterCancelResult", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.core.CompilerParams", "methods": [ { "kind": "function", "name": "CompilerParams.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "CompilerParams.approx_math", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.dimension_semantics", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.max_concurrent_steps", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.unsafe_no_auto_barriers", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.reduction_scratch_bytes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.profile_space", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.profile_dir", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.lowering_semantics", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 8, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.core.TMEMLayout", "methods": [ { "kind": "function", "name": "TMEMLayout.__call__", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "TMEMLayout.to_mgpu", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "TMEMLayout.SCALES_LAYOUT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TMEMLayout.SPARSE_METADATA_LAYOUT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TMEMLayout.M64_COLLECTIVE_LAYOUT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 6, "n_typed": 2, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 4, "n_attrs": 3, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.core.SemaphoreType", "methods": [ { "kind": "function", "name": "SemaphoreType.__call__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SemaphoreType.get_array_aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "SemaphoreType.get_ref_aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "SemaphoreType.REGULAR", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SemaphoreType.BARRIER", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 4, "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 1, "n_attrs": 2, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.core.WarpMesh", "methods": [ { "kind": "function", "name": "WarpMesh.discharges_effect", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "WarpMesh.shape", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "WarpMesh.default_memory_space", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "WarpMesh._NUM_WARPS_PER_WARPGROUP", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "WarpMesh.axis_name", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 4, "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 1, "n_attrs": 2, "n_properties": 2 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.core.BlockSpec", "methods": [ { "kind": "function", "name": "BlockSpec.to_block_mapping", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "BlockSpec.transforms", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSpec.delay_release", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSpec.collective_axes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 8, "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 7, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.core.untile_ref", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [ { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] } ], "name": "jax._src.pallas.mosaic_gpu.core", "names": [ "jax._src.pallas.mosaic_gpu.core.Barrier", "jax._src.pallas.mosaic_gpu.core.BlockSpec", "jax._src.pallas.mosaic_gpu.core.ClusterBarrier", "jax._src.pallas.mosaic_gpu.core.CompilerParams", "jax._src.pallas.mosaic_gpu.core.Layout", "jax._src.pallas.mosaic_gpu.core.MemorySpace", "jax._src.pallas.mosaic_gpu.core.Mesh", "jax._src.pallas.mosaic_gpu.core.PeerMemRef", "jax._src.pallas.mosaic_gpu.core.RefUnion", "jax._src.pallas.mosaic_gpu.core.SemaphoreType", "jax._src.pallas.mosaic_gpu.core.SwizzleTransform", "jax._src.pallas.mosaic_gpu.core.TMEMLayout", "jax._src.pallas.mosaic_gpu.core.TilingTransform", "jax._src.pallas.mosaic_gpu.core.TryClusterCancelResult", "jax._src.pallas.mosaic_gpu.core.WGMMAAccumulatorRef", "jax._src.pallas.mosaic_gpu.core.WarpMesh", "jax._src.pallas.mosaic_gpu.core.kernel", "jax._src.pallas.mosaic_gpu.core.layout_cast", "jax._src.pallas.mosaic_gpu.core.multicast_ref", "jax._src.pallas.mosaic_gpu.core.remote_ref", "jax._src.pallas.mosaic_gpu.core.transform_ref", "jax._src.pallas.mosaic_gpu.core.transpose_ref", "jax._src.pallas.mosaic_gpu.core.unswizzle_ref", "jax._src.pallas.mosaic_gpu.core.untile_ref" ], "n_typable": 106, "n_typed": 77, "n_any": 4, "n_untyped": 25, "n_functions": 9, "n_function_overloads": 9, "n_function_params": 28, "n_methods": 31, "n_method_overloads": 31, "n_method_params": 34, "n_classes": 15, "n_attrs": 60, "n_properties": 4, "n_type_ignores": 3 }, { "path": "jax/_src/pallas/mosaic_gpu/torch.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.torch.as_torch_kernel", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax._src.pallas.mosaic_gpu.torch", "names": [ "jax._src.pallas.mosaic_gpu.torch.as_torch_kernel" ], "n_typable": 2, "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 1, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/mosaic_gpu/primitives.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.wgmma", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.copy_gmem_to_smem", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.multimem_load_reduce", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.wgmma_accumulator_load", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.wait_smem_to_gmem", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.async_store_tmem", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.print_layout", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.semaphore_signal_multicast", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.primitives.SemaphoreSignal", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "SemaphoreSignal.ref", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SemaphoreSignal._", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SemaphoreSignal.device_id", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SemaphoreSignal.inc", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 4, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.primitives.RefType", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "RefType.transforms", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 1, "n_properties": 0 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.query_cluster_cancel", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.commit_tmem", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.tcgen05_mma", "n_typed": 9, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.inline_mgpu", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.load", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.semaphore_signal_parallel", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.set_max_registers", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.barrier_wait", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.commit_smem", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.async_prefetch", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.commit_smem_to_gmem_group", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.try_cluster_cancel", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.multimem_store", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.wgmma_wait", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.async_load_tmem", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.broadcasted_iota", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.wait_load_tmem", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.async_copy_smem_to_tmem", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.tcgen05_commit_arrive", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.async_copy_scales_to_tmem", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.primitives.ShapeDtypeStruct", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "ShapeDtypeStruct.shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ShapeDtypeStruct.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ShapeDtypeStruct.layout", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.barrier_arrive", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.async_copy_sparse_metadata_to_tmem", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.primitives.copy_smem_to_gmem", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "type_ignores": [], "name": "jax._src.pallas.mosaic_gpu.primitives", "names": [ "jax._src.pallas.mosaic_gpu.primitives.RefType", "jax._src.pallas.mosaic_gpu.primitives.SemaphoreSignal", "jax._src.pallas.mosaic_gpu.primitives.ShapeDtypeStruct", "jax._src.pallas.mosaic_gpu.primitives.async_copy_scales_to_tmem", "jax._src.pallas.mosaic_gpu.primitives.async_copy_smem_to_tmem", "jax._src.pallas.mosaic_gpu.primitives.async_copy_sparse_metadata_to_tmem", "jax._src.pallas.mosaic_gpu.primitives.async_load_tmem", "jax._src.pallas.mosaic_gpu.primitives.async_prefetch", "jax._src.pallas.mosaic_gpu.primitives.async_store_tmem", "jax._src.pallas.mosaic_gpu.primitives.barrier_arrive", "jax._src.pallas.mosaic_gpu.primitives.barrier_wait", "jax._src.pallas.mosaic_gpu.primitives.broadcasted_iota", "jax._src.pallas.mosaic_gpu.primitives.commit_smem", "jax._src.pallas.mosaic_gpu.primitives.commit_smem_to_gmem_group", "jax._src.pallas.mosaic_gpu.primitives.commit_tmem", "jax._src.pallas.mosaic_gpu.primitives.copy_gmem_to_smem", "jax._src.pallas.mosaic_gpu.primitives.copy_smem_to_gmem", "jax._src.pallas.mosaic_gpu.primitives.inline_mgpu", "jax._src.pallas.mosaic_gpu.primitives.load", "jax._src.pallas.mosaic_gpu.primitives.multimem_load_reduce", "jax._src.pallas.mosaic_gpu.primitives.multimem_store", "jax._src.pallas.mosaic_gpu.primitives.print_layout", "jax._src.pallas.mosaic_gpu.primitives.query_cluster_cancel", "jax._src.pallas.mosaic_gpu.primitives.semaphore_signal_multicast", "jax._src.pallas.mosaic_gpu.primitives.semaphore_signal_parallel", "jax._src.pallas.mosaic_gpu.primitives.set_max_registers", "jax._src.pallas.mosaic_gpu.primitives.tcgen05_commit_arrive", "jax._src.pallas.mosaic_gpu.primitives.tcgen05_mma", "jax._src.pallas.mosaic_gpu.primitives.try_cluster_cancel", "jax._src.pallas.mosaic_gpu.primitives.wait_load_tmem", "jax._src.pallas.mosaic_gpu.primitives.wait_smem_to_gmem", "jax._src.pallas.mosaic_gpu.primitives.wgmma", "jax._src.pallas.mosaic_gpu.primitives.wgmma_accumulator_load", "jax._src.pallas.mosaic_gpu.primitives.wgmma_wait" ], "n_typable": 106, "n_typed": 82, "n_any": 0, "n_untyped": 24, "n_functions": 31, "n_function_overloads": 31, "n_function_params": 75, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 3, "n_attrs": 8, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/mosaic_gpu/helpers.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.helpers.find_swizzle", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.helpers.dynamic_scheduling_loop", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 2, "n_typable": 4, "n_params": 3, "n_function_overloads": 2, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.helpers.planar_snake", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.helpers.nd_loop", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 2, "n_typable": 5, "n_params": 4, "n_function_overloads": 2, "n_function_params": 4 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.helpers.NDLoopInfo", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "NDLoopInfo.index", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "NDLoopInfo.local_index", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "NDLoopInfo.num_local_steps", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.helpers.format_tcgen05_sparse_metadata", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax._src.pallas.mosaic_gpu.helpers", "names": [ "jax._src.pallas.mosaic_gpu.helpers.NDLoopInfo", "jax._src.pallas.mosaic_gpu.helpers.dynamic_scheduling_loop", "jax._src.pallas.mosaic_gpu.helpers.find_swizzle", "jax._src.pallas.mosaic_gpu.helpers.format_tcgen05_sparse_metadata", "jax._src.pallas.mosaic_gpu.helpers.nd_loop", "jax._src.pallas.mosaic_gpu.helpers.planar_snake" ], "n_typable": 19, "n_typed": 15, "n_any": 0, "n_untyped": 4, "n_functions": 5, "n_function_overloads": 7, "n_function_params": 14, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 3, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/mosaic_gpu/pipeline.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.pipeline.emit_pipeline", "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "class", "name": "jax._src.pallas.mosaic_gpu.pipeline.PipelinePipeline", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "PipelinePipeline.START", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "PipelinePipeline.STEADY", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "PipelinePipeline.STOP", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax._src.pallas.mosaic_gpu.pipeline.emit_pipeline_warp_specialized", "n_typed": 13, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 13, "n_params": 12, "n_function_overloads": 1, "n_function_params": 12 } ], "type_ignores": [], "name": "jax._src.pallas.mosaic_gpu.pipeline", "names": [ "jax._src.pallas.mosaic_gpu.pipeline.PipelinePipeline", "jax._src.pallas.mosaic_gpu.pipeline.emit_pipeline", "jax._src.pallas.mosaic_gpu.pipeline.emit_pipeline_warp_specialized" ], "n_typable": 20, "n_typed": 19, "n_any": 0, "n_untyped": 1, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 18, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 3, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/gpu/all_gather_mgpu.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.all_gather_mgpu.all_gather", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.gpu.all_gather_mgpu", "names": [ "jax.experimental.pallas.ops.gpu.all_gather_mgpu.all_gather" ], "n_typable": 7, "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 6, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/gpu/attention.py", "symbol_reports": [ { "kind": "class", "name": "jax.experimental.pallas.ops.gpu.attention.BlockSizes", "methods": [ { "kind": "function", "name": "BlockSizes.get_default", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "BlockSizes.has_backward_blocks", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "BlockSizes.block_q", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_k", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_q_dkv", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_kv_dkv", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_q_dq", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_kv_dq", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 2, "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 6, "n_properties": 1 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.gpu.attention.DEFAULT_MASK_VALUE", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.attention.segment_mask", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.attention.mha", "n_typed": 11, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 15, "n_params": 14, "n_function_overloads": 1, "n_function_params": 14 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.attention.mha_reference", "n_typed": 2, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.attention.mha_forward_kernel", "n_typed": 6, "n_any": 2, "n_untyped": 4, "n_overloads": 1, "n_typable": 12, "n_params": 11, "n_function_overloads": 1, "n_function_params": 11 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.attention.mha_backward_kernel", "n_typed": 8, "n_any": 0, "n_untyped": 11, "n_overloads": 1, "n_typable": 19, "n_params": 18, "n_function_overloads": 1, "n_function_params": 18 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.gpu.attention", "names": [ "jax.experimental.pallas.ops.gpu.attention.BlockSizes", "jax.experimental.pallas.ops.gpu.attention.DEFAULT_MASK_VALUE", "jax.experimental.pallas.ops.gpu.attention.mha", "jax.experimental.pallas.ops.gpu.attention.mha_backward_kernel", "jax.experimental.pallas.ops.gpu.attention.mha_forward_kernel", "jax.experimental.pallas.ops.gpu.attention.mha_reference", "jax.experimental.pallas.ops.gpu.attention.segment_mask" ], "n_typable": 59, "n_typed": 30, "n_any": 2, "n_untyped": 27, "n_functions": 5, "n_function_overloads": 5, "n_function_params": 51, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_classes": 1, "n_attrs": 7, "n_properties": 1, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/gpu/attention_mgpu.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.attention_mgpu.attention_with_pipeline_emitter", "n_typed": 1, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.attention_mgpu.main", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.attention_mgpu.attention_reference", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "class", "name": "jax.experimental.pallas.ops.gpu.attention_mgpu.TuningConfig", "methods": [ { "kind": "function", "name": "TuningConfig.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "TuningConfig.has_backward_blocks", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "TuningConfig.block_q", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.block_kv", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.max_concurrent_steps", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.use_schedule_barrier", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.causal", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.compute_wgs_bwd", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.block_q_dkv", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.block_kv_dkv", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.block_q_dq", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.block_kv_dq", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 2, "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 10, "n_properties": 1 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.attention_mgpu.attention", "n_typed": 2, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.gpu.attention_mgpu", "names": [ "jax.experimental.pallas.ops.gpu.attention_mgpu.TuningConfig", "jax.experimental.pallas.ops.gpu.attention_mgpu.attention", "jax.experimental.pallas.ops.gpu.attention_mgpu.attention_reference", "jax.experimental.pallas.ops.gpu.attention_mgpu.attention_with_pipeline_emitter", "jax.experimental.pallas.ops.gpu.attention_mgpu.main" ], "n_typable": 22, "n_typed": 4, "n_any": 0, "n_untyped": 18, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 16, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_classes": 1, "n_attrs": 10, "n_properties": 1, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/gpu/blackwell_matmul_mgpu.py", "symbol_reports": [ { "kind": "class", "name": "jax.experimental.pallas.ops.gpu.blackwell_matmul_mgpu.TuningConfig", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "TuningConfig.tile_m", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.tile_n", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.tile_k", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.max_concurrent_steps", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.collective", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.epilogue_tile_n", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.grid_minor_dim", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.grid_tile_width", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 8, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.pallas.ops.gpu.blackwell_matmul_mgpu.MatmulDimension", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "MatmulDimension.M", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MatmulDimension.N", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.blackwell_matmul_mgpu.main", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.blackwell_matmul_mgpu.matmul_kernel", "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.gpu.blackwell_matmul_mgpu", "names": [ "jax.experimental.pallas.ops.gpu.blackwell_matmul_mgpu.MatmulDimension", "jax.experimental.pallas.ops.gpu.blackwell_matmul_mgpu.TuningConfig", "jax.experimental.pallas.ops.gpu.blackwell_matmul_mgpu.main", "jax.experimental.pallas.ops.gpu.blackwell_matmul_mgpu.matmul_kernel" ], "n_typable": 6, "n_typed": 2, "n_any": 0, "n_untyped": 4, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 4, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 2, "n_attrs": 10, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/gpu/blackwell_ragged_dot_mgpu.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.blackwell_ragged_dot_mgpu.main", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.blackwell_ragged_dot_mgpu.ragged_dot_kernel", "n_typed": 1, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.blackwell_ragged_dot_mgpu.do_matmul", "n_typed": 6, "n_any": 0, "n_untyped": 13, "n_overloads": 1, "n_typable": 19, "n_params": 18, "n_function_overloads": 1, "n_function_params": 18 }, { "kind": "class", "name": "jax.experimental.pallas.ops.gpu.blackwell_ragged_dot_mgpu.TuningConfig", "methods": [ { "kind": "function", "name": "TuningConfig.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "TuningConfig.tile_m", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.tile_n", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.tile_k", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.max_concurrent_steps", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.collective", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.grid_tile_width", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.grid_minor_dim", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.epilogue_tile_n", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 8, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.blackwell_ragged_dot_mgpu.sample_group_sizes", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.blackwell_ragged_dot_mgpu.ragged_dot_reference", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.gpu.blackwell_ragged_dot_mgpu", "names": [ "jax.experimental.pallas.ops.gpu.blackwell_ragged_dot_mgpu.TuningConfig", "jax.experimental.pallas.ops.gpu.blackwell_ragged_dot_mgpu.do_matmul", "jax.experimental.pallas.ops.gpu.blackwell_ragged_dot_mgpu.main", "jax.experimental.pallas.ops.gpu.blackwell_ragged_dot_mgpu.ragged_dot_kernel", "jax.experimental.pallas.ops.gpu.blackwell_ragged_dot_mgpu.ragged_dot_reference", "jax.experimental.pallas.ops.gpu.blackwell_ragged_dot_mgpu.sample_group_sizes" ], "n_typable": 36, "n_typed": 12, "n_any": 0, "n_untyped": 24, "n_functions": 5, "n_function_overloads": 5, "n_function_params": 30, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_classes": 1, "n_attrs": 8, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/gpu/collective_matmul_mgpu.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.collective_matmul_mgpu.is_nvshmem_used", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.collective_matmul_mgpu.all_gather_lhs_matmul", "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.gpu.collective_matmul_mgpu", "names": [ "jax.experimental.pallas.ops.gpu.collective_matmul_mgpu.all_gather_lhs_matmul", "jax.experimental.pallas.ops.gpu.collective_matmul_mgpu.is_nvshmem_used" ], "n_typable": 7, "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 5, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/gpu/decode_attention.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.decode_attention.gqa", "n_typed": 11, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 17, "n_params": 16, "n_function_overloads": 1, "n_function_params": 16 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.decode_attention.gqa_reference", "n_typed": 0, "n_any": 0, "n_untyped": 9, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.decode_attention.mqa", "n_typed": 11, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 17, "n_params": 16, "n_function_overloads": 1, "n_function_params": 16 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.decode_attention.attn_forward_kernel", "n_typed": 4, "n_any": 2, "n_untyped": 6, "n_overloads": 1, "n_typable": 12, "n_params": 11, "n_function_overloads": 1, "n_function_params": 11 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.decode_attention.mqa_reference", "n_typed": 0, "n_any": 0, "n_untyped": 9, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.decode_attention.mha_reference", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.decode_attention.decode_attn_unbatched", "n_typed": 11, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 17, "n_params": 16, "n_function_overloads": 1, "n_function_params": 16 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.gpu.decode_attention", "names": [ "jax.experimental.pallas.ops.gpu.decode_attention.attn_forward_kernel", "jax.experimental.pallas.ops.gpu.decode_attention.decode_attn_unbatched", "jax.experimental.pallas.ops.gpu.decode_attention.gqa", "jax.experimental.pallas.ops.gpu.decode_attention.gqa_reference", "jax.experimental.pallas.ops.gpu.decode_attention.mha_reference", "jax.experimental.pallas.ops.gpu.decode_attention.mqa", "jax.experimental.pallas.ops.gpu.decode_attention.mqa_reference" ], "n_typable": 88, "n_typed": 37, "n_any": 2, "n_untyped": 49, "n_functions": 7, "n_function_overloads": 7, "n_function_params": 81, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/gpu/hopper_matmul_mgpu.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.hopper_matmul_mgpu.kernel", "n_typed": 1, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.hopper_matmul_mgpu.main", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax.experimental.pallas.ops.gpu.hopper_matmul_mgpu.MatmulDimension", "methods": [ { "kind": "function", "name": "MatmulDimension.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "MatmulDimension.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "MatmulDimension.M", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MatmulDimension.N", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 2, "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.hopper_matmul_mgpu.matmul", "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax.experimental.pallas.ops.gpu.hopper_matmul_mgpu.TuningConfig", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "TuningConfig.tile_m", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.tile_n", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.tile_k", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.max_concurrent_steps", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.epi_tile_n", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.epi_tile_m", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.grid_minor_dim", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.grid_tile_width", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.wg_dimension", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.cluster_dimension", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 10, "n_properties": 0 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.gpu.hopper_matmul_mgpu", "names": [ "jax.experimental.pallas.ops.gpu.hopper_matmul_mgpu.MatmulDimension", "jax.experimental.pallas.ops.gpu.hopper_matmul_mgpu.TuningConfig", "jax.experimental.pallas.ops.gpu.hopper_matmul_mgpu.kernel", "jax.experimental.pallas.ops.gpu.hopper_matmul_mgpu.main", "jax.experimental.pallas.ops.gpu.hopper_matmul_mgpu.matmul" ], "n_typable": 15, "n_typed": 3, "n_any": 0, "n_untyped": 12, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 10, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 0, "n_classes": 2, "n_attrs": 12, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/gpu/hopper_mixed_type_matmul_mgpu.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.hopper_mixed_type_matmul_mgpu.reference", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.hopper_mixed_type_matmul_mgpu.main", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax.experimental.pallas.ops.gpu.hopper_mixed_type_matmul_mgpu.MatmulDimension", "methods": [ { "kind": "function", "name": "MatmulDimension.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "MatmulDimension.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "MatmulDimension.M", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MatmulDimension.N", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 2, "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.hopper_mixed_type_matmul_mgpu.mixed_matmul_kernel", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "class", "name": "jax.experimental.pallas.ops.gpu.hopper_mixed_type_matmul_mgpu.TuningConfig", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "TuningConfig.tile_m", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.tile_n", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.tile_k", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.max_concurrent_steps", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.epi_tile_n", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.epi_tile_m", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.grid_minor_dim", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.grid_tile_width", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.wg_dimension", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TuningConfig.cluster_dimension", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 10, "n_properties": 0 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.gpu.hopper_mixed_type_matmul_mgpu", "names": [ "jax.experimental.pallas.ops.gpu.hopper_mixed_type_matmul_mgpu.MatmulDimension", "jax.experimental.pallas.ops.gpu.hopper_mixed_type_matmul_mgpu.TuningConfig", "jax.experimental.pallas.ops.gpu.hopper_mixed_type_matmul_mgpu.main", "jax.experimental.pallas.ops.gpu.hopper_mixed_type_matmul_mgpu.mixed_matmul_kernel", "jax.experimental.pallas.ops.gpu.hopper_mixed_type_matmul_mgpu.reference" ], "n_typable": 13, "n_typed": 10, "n_any": 0, "n_untyped": 3, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 8, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 0, "n_classes": 2, "n_attrs": 12, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/gpu/layer_norm.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.layer_norm.layer_norm_reference", "n_typed": 1, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.layer_norm.layer_norm_forward_kernel", "n_typed": 2, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.layer_norm.layer_norm_backward_kernel_dx", "n_typed": 2, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.layer_norm.layer_norm", "n_typed": 5, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.layer_norm.layer_norm_backward_kernel_dw_db", "n_typed": 3, "n_any": 0, "n_untyped": 9, "n_overloads": 1, "n_typable": 12, "n_params": 11, "n_function_overloads": 1, "n_function_params": 11 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.layer_norm.layer_norm_backward", "n_typed": 5, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.layer_norm.layer_norm_forward", "n_typed": 5, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.gpu.layer_norm", "names": [ "jax.experimental.pallas.ops.gpu.layer_norm.layer_norm", "jax.experimental.pallas.ops.gpu.layer_norm.layer_norm_backward", "jax.experimental.pallas.ops.gpu.layer_norm.layer_norm_backward_kernel_dw_db", "jax.experimental.pallas.ops.gpu.layer_norm.layer_norm_backward_kernel_dx", "jax.experimental.pallas.ops.gpu.layer_norm.layer_norm_forward", "jax.experimental.pallas.ops.gpu.layer_norm.layer_norm_forward_kernel", "jax.experimental.pallas.ops.gpu.layer_norm.layer_norm_reference" ], "n_typable": 62, "n_typed": 23, "n_any": 0, "n_untyped": 39, "n_functions": 7, "n_function_overloads": 7, "n_function_params": 55, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/gpu/paged_attention.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.paged_attention.paged_attention_reference", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.gpu.paged_attention.DEFAULT_MASK_VALUE", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.paged_attention.paged_attention_kernel", "n_typed": 4, "n_any": 2, "n_untyped": 8, "n_overloads": 1, "n_typable": 14, "n_params": 13, "n_function_overloads": 1, "n_function_params": 13 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.paged_attention.paged_attention_unbatched", "n_typed": 17, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 17, "n_params": 16, "n_function_overloads": 1, "n_function_params": 16 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.paged_attention.paged_attention", "n_typed": 17, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 17, "n_params": 16, "n_function_overloads": 1, "n_function_params": 16 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.gpu.paged_attention", "names": [ "jax.experimental.pallas.ops.gpu.paged_attention.DEFAULT_MASK_VALUE", "jax.experimental.pallas.ops.gpu.paged_attention.paged_attention", "jax.experimental.pallas.ops.gpu.paged_attention.paged_attention_kernel", "jax.experimental.pallas.ops.gpu.paged_attention.paged_attention_reference", "jax.experimental.pallas.ops.gpu.paged_attention.paged_attention_unbatched" ], "n_typable": 56, "n_typed": 45, "n_any": 2, "n_untyped": 9, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 51, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/gpu/ragged_dot_mgpu.py", "symbol_reports": [ { "kind": "class", "name": "jax.experimental.pallas.ops.gpu.ragged_dot_mgpu.GroupInfo", "methods": [ { "kind": "function", "name": "GroupInfo.create", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "GroupInfo.group_id", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GroupInfo.block", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GroupInfo.block_start", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GroupInfo.actual_start", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GroupInfo.actual_end", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GroupInfo.start_within_block", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GroupInfo.actual_size", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 4, "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 3, "n_attrs": 7, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.ragged_dot_mgpu.ragged_dot", "n_typed": 8, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 11, "n_params": 10, "n_function_overloads": 1, "n_function_params": 10 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.ragged_dot_mgpu.main", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.gpu.ragged_dot_mgpu", "names": [ "jax.experimental.pallas.ops.gpu.ragged_dot_mgpu.GroupInfo", "jax.experimental.pallas.ops.gpu.ragged_dot_mgpu.main", "jax.experimental.pallas.ops.gpu.ragged_dot_mgpu.ragged_dot" ], "n_typable": 17, "n_typed": 8, "n_any": 0, "n_untyped": 9, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 11, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 3, "n_classes": 1, "n_attrs": 7, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/gpu/reduce_scatter_mgpu.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.reduce_scatter_mgpu.reduce_scatter", "n_typed": 7, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.gpu.reduce_scatter_mgpu", "names": [ "jax.experimental.pallas.ops.gpu.reduce_scatter_mgpu.reduce_scatter" ], "n_typable": 8, "n_typed": 7, "n_any": 0, "n_untyped": 1, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 7, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/gpu/rms_norm.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.rms_norm.rms_norm", "n_typed": 5, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.rms_norm.rms_norm_forward", "n_typed": 5, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.rms_norm.rms_norm_backward_kernel_dx", "n_typed": 2, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.rms_norm.rms_norm_backward", "n_typed": 5, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.rms_norm.rms_norm_backward_kernel_dw_db", "n_typed": 3, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 11, "n_params": 10, "n_function_overloads": 1, "n_function_params": 10 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.rms_norm.rms_norm_reference", "n_typed": 1, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.rms_norm.rms_norm_forward_kernel", "n_typed": 2, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.gpu.rms_norm", "names": [ "jax.experimental.pallas.ops.gpu.rms_norm.rms_norm", "jax.experimental.pallas.ops.gpu.rms_norm.rms_norm_backward", "jax.experimental.pallas.ops.gpu.rms_norm.rms_norm_backward_kernel_dw_db", "jax.experimental.pallas.ops.gpu.rms_norm.rms_norm_backward_kernel_dx", "jax.experimental.pallas.ops.gpu.rms_norm.rms_norm_forward", "jax.experimental.pallas.ops.gpu.rms_norm.rms_norm_forward_kernel", "jax.experimental.pallas.ops.gpu.rms_norm.rms_norm_reference" ], "n_typable": 59, "n_typed": 23, "n_any": 0, "n_untyped": 36, "n_functions": 7, "n_function_overloads": 7, "n_function_params": 52, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/gpu/softmax.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.softmax.softmax", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.gpu.softmax", "names": [ "jax.experimental.pallas.ops.gpu.softmax.softmax" ], "n_typable": 6, "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 5, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/gpu/transposed_ragged_dot_mgpu.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.transposed_ragged_dot_mgpu.ref_transposed_ragged_dot", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.transposed_ragged_dot_mgpu.transposed_ragged_dot", "n_typed": 6, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.experimental.pallas.ops.gpu.transposed_ragged_dot_mgpu.main", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.gpu.transposed_ragged_dot_mgpu", "names": [ "jax.experimental.pallas.ops.gpu.transposed_ragged_dot_mgpu.main", "jax.experimental.pallas.ops.gpu.transposed_ragged_dot_mgpu.ref_transposed_ragged_dot", "jax.experimental.pallas.ops.gpu.transposed_ragged_dot_mgpu.transposed_ragged_dot" ], "n_typable": 15, "n_typed": 6, "n_any": 0, "n_untyped": 9, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 12, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/tpu/all_gather.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.all_gather.ag_kernel", "n_typed": 2, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.all_gather.all_gather", "n_typed": 3, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.all_gather.get_neighbor", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.tpu.all_gather", "names": [ "jax.experimental.pallas.ops.tpu.all_gather.ag_kernel", "jax.experimental.pallas.ops.tpu.all_gather.all_gather", "jax.experimental.pallas.ops.tpu.all_gather.get_neighbor" ], "n_typable": 17, "n_typed": 10, "n_any": 0, "n_untyped": 7, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 14, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/tpu/example_kernel.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.example_kernel.double_kernel", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.example_kernel.double", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.tpu.example_kernel", "names": [ "jax.experimental.pallas.ops.tpu.example_kernel.double", "jax.experimental.pallas.ops.tpu.example_kernel.double_kernel" ], "n_typable": 5, "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 3, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/tpu/flash_attention.py", "symbol_reports": [ { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.flash_attention.BlockSizes", "methods": [ { "kind": "function", "name": "BlockSizes.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BlockSizes.get_default", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "properties": [ { "kind": "property", "name": "BlockSizes.has_backward_blocks", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "BlockSizes.block_q", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_k_major", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_k", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_b", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_q_major_dkv", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_k_major_dkv", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_k_dkv", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_q_dkv", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_k_major_dq", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_k_dq", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_q_dq", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 8, "n_typed": 1, "n_any": 0, "n_untyped": 7, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 5, "n_attrs": 11, "n_properties": 1 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.flash_attention.DEFAULT_MASK_VALUE", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.flash_attention.mha_reference_no_custom_vjp", "n_typed": 6, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.flash_attention.below_or_on_diag", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.flash_attention.NUM_SUBLANES", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.flash_attention.NUM_LANES", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.flash_attention.MIN_BLOCK_SIZE", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.flash_attention.TRANS_B_DIM_NUMBERS", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.flash_attention.mha_reference_bwd", "n_typed": 4, "n_any": 0, "n_untyped": 9, "n_overloads": 1, "n_typable": 13, "n_params": 12, "n_function_overloads": 1, "n_function_params": 12 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.flash_attention.mha_reference", "n_typed": 3, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.flash_attention.SegmentIds", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "SegmentIds.q", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SegmentIds.kv", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.flash_attention.flash_attention", "n_typed": 4, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.tpu.flash_attention", "names": [ "jax.experimental.pallas.ops.tpu.flash_attention.BlockSizes", "jax.experimental.pallas.ops.tpu.flash_attention.DEFAULT_MASK_VALUE", "jax.experimental.pallas.ops.tpu.flash_attention.MIN_BLOCK_SIZE", "jax.experimental.pallas.ops.tpu.flash_attention.NUM_LANES", "jax.experimental.pallas.ops.tpu.flash_attention.NUM_SUBLANES", "jax.experimental.pallas.ops.tpu.flash_attention.SegmentIds", "jax.experimental.pallas.ops.tpu.flash_attention.TRANS_B_DIM_NUMBERS", "jax.experimental.pallas.ops.tpu.flash_attention.below_or_on_diag", "jax.experimental.pallas.ops.tpu.flash_attention.flash_attention", "jax.experimental.pallas.ops.tpu.flash_attention.mha_reference", "jax.experimental.pallas.ops.tpu.flash_attention.mha_reference_bwd", "jax.experimental.pallas.ops.tpu.flash_attention.mha_reference_no_custom_vjp" ], "n_typable": 56, "n_typed": 18, "n_any": 0, "n_untyped": 38, "n_functions": 5, "n_function_overloads": 5, "n_function_params": 42, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 5, "n_classes": 2, "n_attrs": 18, "n_properties": 1, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/tpu/matmul.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.matmul.matmul", "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.matmul.matmul_kernel", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.tpu.matmul", "names": [ "jax.experimental.pallas.ops.tpu.matmul.matmul", "jax.experimental.pallas.ops.tpu.matmul.matmul_kernel" ], "n_typable": 12, "n_typed": 6, "n_any": 0, "n_untyped": 6, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 10, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/tpu/megablox/ops.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.megablox.ops.gmm", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.tpu.megablox.ops", "names": [ "jax.experimental.pallas.ops.tpu.megablox.ops.gmm" ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/tpu/megablox/common.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.megablox.common.select_input_dtype", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.megablox.common.tpu_kind", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.megablox.common.supports_bfloat16_matmul", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.megablox.common.tpu_generation", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.megablox.common.assert_is_supported_dtype", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.megablox.common.is_tpu", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.tpu.megablox.common", "names": [ "jax.experimental.pallas.ops.tpu.megablox.common.assert_is_supported_dtype", "jax.experimental.pallas.ops.tpu.megablox.common.is_tpu", "jax.experimental.pallas.ops.tpu.megablox.common.select_input_dtype", "jax.experimental.pallas.ops.tpu.megablox.common.supports_bfloat16_matmul", "jax.experimental.pallas.ops.tpu.megablox.common.tpu_generation", "jax.experimental.pallas.ops.tpu.megablox.common.tpu_kind" ], "n_typable": 9, "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_functions": 6, "n_function_overloads": 6, "n_function_params": 3, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/tpu/megablox/gmm.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.megablox.gmm.tgmm", "n_typed": 10, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.megablox.gmm.gmm", "n_typed": 10, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.megablox.gmm.LutFn", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.megablox.gmm.make_group_metadata", "n_typed": 6, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.tpu.megablox.gmm", "names": [ "jax.experimental.pallas.ops.tpu.megablox.gmm.LutFn", "jax.experimental.pallas.ops.tpu.megablox.gmm.gmm", "jax.experimental.pallas.ops.tpu.megablox.gmm.make_group_metadata", "jax.experimental.pallas.ops.tpu.megablox.gmm.tgmm" ], "n_typable": 28, "n_typed": 27, "n_any": 1, "n_untyped": 0, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 24, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.paged_attention.paged_attention_kernel.paged_attention", "n_typed": 11, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 11, "n_params": 10, "n_function_overloads": 1, "n_function_params": 10 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.paged_attention.paged_attention_kernel.DEFAULT_MASK_VALUE", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.paged_attention.paged_attention_kernel.paged_flash_attention_kernel", "n_typed": 6, "n_any": 0, "n_untyped": 20, "n_overloads": 1, "n_typable": 26, "n_params": 25, "n_function_overloads": 1, "n_function_params": 25 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.paged_attention.paged_attention_kernel.paged_flash_attention_kernel_inline_seq_dim", "n_typed": 6, "n_any": 0, "n_untyped": 19, "n_overloads": 1, "n_typable": 25, "n_params": 24, "n_function_overloads": 1, "n_function_params": 24 }, { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.paged_attention.paged_attention_kernel.MultiPageAsyncCopyDescriptor", "methods": [ { "kind": "function", "name": "MultiPageAsyncCopyDescriptor.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 10, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "MultiPageAsyncCopyDescriptor._make_async_copy", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "MultiPageAsyncCopyDescriptor._make_scales_async_copy", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "MultiPageAsyncCopyDescriptor.start", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "MultiPageAsyncCopyDescriptor._maybe_dequantize", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "MultiPageAsyncCopyDescriptor.wait_and_get_loaded", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [], "n_typable": 20, "n_typed": 1, "n_any": 0, "n_untyped": 19, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 14, "n_attrs": 0, "n_properties": 0 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.tpu.paged_attention.paged_attention_kernel", "names": [ "jax.experimental.pallas.ops.tpu.paged_attention.paged_attention_kernel.DEFAULT_MASK_VALUE", "jax.experimental.pallas.ops.tpu.paged_attention.paged_attention_kernel.MultiPageAsyncCopyDescriptor", "jax.experimental.pallas.ops.tpu.paged_attention.paged_attention_kernel.paged_attention", "jax.experimental.pallas.ops.tpu.paged_attention.paged_attention_kernel.paged_flash_attention_kernel", "jax.experimental.pallas.ops.tpu.paged_attention.paged_attention_kernel.paged_flash_attention_kernel_inline_seq_dim" ], "n_typable": 83, "n_typed": 24, "n_any": 0, "n_untyped": 59, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 59, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 14, "n_classes": 1, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/tpu/paged_attention/quantization_utils.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.paged_attention.quantization_utils.to_int8", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.paged_attention.quantization_utils.MAX_INT8", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.paged_attention.quantization_utils.QuantizedTensor", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "QuantizedTensor.weight", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "QuantizedTensor.scales", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.paged_attention.quantization_utils.from_int8", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.paged_attention.quantization_utils.get_quantization_scales", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.paged_attention.quantization_utils.unquantize_from_int8", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.paged_attention.quantization_utils.quantize_to_int8", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.tpu.paged_attention.quantization_utils", "names": [ "jax.experimental.pallas.ops.tpu.paged_attention.quantization_utils.MAX_INT8", "jax.experimental.pallas.ops.tpu.paged_attention.quantization_utils.QuantizedTensor", "jax.experimental.pallas.ops.tpu.paged_attention.quantization_utils.from_int8", "jax.experimental.pallas.ops.tpu.paged_attention.quantization_utils.get_quantization_scales", "jax.experimental.pallas.ops.tpu.paged_attention.quantization_utils.quantize_to_int8", "jax.experimental.pallas.ops.tpu.paged_attention.quantization_utils.to_int8", "jax.experimental.pallas.ops.tpu.paged_attention.quantization_utils.unquantize_from_int8" ], "n_typable": 14, "n_typed": 14, "n_any": 0, "n_untyped": 0, "n_functions": 5, "n_function_overloads": 5, "n_function_params": 9, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 3, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/tpu/paged_attention/util.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.paged_attention.util.grouped_query_attention_reference", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.paged_attention.util.MASK_VALUE", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.tpu.paged_attention.util", "names": [ "jax.experimental.pallas.ops.tpu.paged_attention.util.MASK_VALUE", "jax.experimental.pallas.ops.tpu.paged_attention.util.grouped_query_attention_reference" ], "n_typable": 8, "n_typed": 7, "n_any": 0, "n_untyped": 1, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 6, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel.ref_ragged_paged_attention", "n_typed": 12, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 13, "n_params": 12, "n_function_overloads": 1, "n_function_params": 12 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel.ragged_paged_attention", "n_typed": 15, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 16, "n_params": 15, "n_function_overloads": 1, "n_function_params": 15 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel.DEFAULT_MASK_VALUE", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel.dynamic_validate_inputs", "n_typed": 15, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 16, "n_params": 15, "n_function_overloads": 1, "n_function_params": 15 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel.get_dtype_packing", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel.static_validate_inputs", "n_typed": 15, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 16, "n_params": 15, "n_function_overloads": 1, "n_function_params": 15 }, { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel.MultiPageAsyncCopyDescriptor", "methods": [ { "kind": "function", "name": "MultiPageAsyncCopyDescriptor.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "MultiPageAsyncCopyDescriptor.start", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "MultiPageAsyncCopyDescriptor.wait", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [], "n_typable": 8, "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 5, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel.ragged_paged_attention_kernel", "n_typed": 6, "n_any": 0, "n_untyped": 14, "n_overloads": 1, "n_typable": 20, "n_params": 19, "n_function_overloads": 1, "n_function_params": 19 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel.get_min_heads_per_blk", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel", "names": [ "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel.DEFAULT_MASK_VALUE", "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel.MultiPageAsyncCopyDescriptor", "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel.dynamic_validate_inputs", "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel.get_dtype_packing", "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel.get_min_heads_per_blk", "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel.ragged_paged_attention", "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel.ragged_paged_attention_kernel", "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel.ref_ragged_paged_attention", "jax.experimental.pallas.ops.tpu.ragged_paged_attention.kernel.static_validate_inputs" ], "n_typable": 97, "n_typed": 63, "n_any": 0, "n_untyped": 34, "n_functions": 7, "n_function_overloads": 7, "n_function_params": 81, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 5, "n_classes": 1, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/tpu/ragged_paged_attention/tuned_block_sizes.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes.get_device_name", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes.MAX_PAGES_PER_SEQ", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes.get_min_page_size", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes.get_tuned_block_sizes", "n_typed": 1, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes.TUNED_BLOCK_SIZES", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes.simplify_key", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes.get_tpu_version", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes.next_power_of_2", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes", "names": [ "jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes.MAX_PAGES_PER_SEQ", "jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes.TUNED_BLOCK_SIZES", "jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes.get_device_name", "jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes.get_min_page_size", "jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes.get_tpu_version", "jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes.get_tuned_block_sizes", "jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes.next_power_of_2", "jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes.simplify_key" ], "n_typable": 19, "n_typed": 4, "n_any": 0, "n_untyped": 15, "n_functions": 6, "n_function_overloads": 6, "n_function_params": 13, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 2, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/tpu/random/philox.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.random.philox.K_HI_32", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.random.philox.MUL_B", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.random.philox.K_LO_32", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.random.philox.BLOCK_SIZE", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.random.philox.philox_4x32", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.random.philox.philox_split", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.random.philox.philox_fold_in", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.random.philox.philox_4x32_count", "n_typed": 3, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.random.philox.mul32_hi_lo", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.random.philox.philox_4x32_kernel", "n_typed": 5, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.random.philox.Shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.random.philox.philox_random_bits", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.random.philox.plphilox_prng_impl", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.random.philox.MUL_A", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.tpu.random.philox", "names": [ "jax.experimental.pallas.ops.tpu.random.philox.BLOCK_SIZE", "jax.experimental.pallas.ops.tpu.random.philox.K_HI_32", "jax.experimental.pallas.ops.tpu.random.philox.K_LO_32", "jax.experimental.pallas.ops.tpu.random.philox.MUL_A", "jax.experimental.pallas.ops.tpu.random.philox.MUL_B", "jax.experimental.pallas.ops.tpu.random.philox.Shape", "jax.experimental.pallas.ops.tpu.random.philox.mul32_hi_lo", "jax.experimental.pallas.ops.tpu.random.philox.philox_4x32", "jax.experimental.pallas.ops.tpu.random.philox.philox_4x32_count", "jax.experimental.pallas.ops.tpu.random.philox.philox_4x32_kernel", "jax.experimental.pallas.ops.tpu.random.philox.philox_fold_in", "jax.experimental.pallas.ops.tpu.random.philox.philox_random_bits", "jax.experimental.pallas.ops.tpu.random.philox.philox_split", "jax.experimental.pallas.ops.tpu.random.philox.plphilox_prng_impl" ], "n_typable": 35, "n_typed": 15, "n_any": 0, "n_untyped": 20, "n_functions": 7, "n_function_overloads": 7, "n_function_params": 26, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 7, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/tpu/random/prng_utils.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.random.prng_utils.compute_scalar_offset", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.random.prng_utils.blocked_iota", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.random.prng_utils.round_up", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.random.prng_utils.Shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.tpu.random.prng_utils", "names": [ "jax.experimental.pallas.ops.tpu.random.prng_utils.Shape", "jax.experimental.pallas.ops.tpu.random.prng_utils.blocked_iota", "jax.experimental.pallas.ops.tpu.random.prng_utils.compute_scalar_offset", "jax.experimental.pallas.ops.tpu.random.prng_utils.round_up" ], "n_typable": 8, "n_typed": 5, "n_any": 0, "n_untyped": 3, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 5, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 2, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/tpu/random/threefry.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.random.threefry.BLOCK_SIZE", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.random.threefry.plthreefry_random_bits", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.random.threefry.threefry_2x32_count", "n_typed": 3, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.random.threefry.plthreefry_prng_impl", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.random.threefry.Shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.tpu.random.threefry", "names": [ "jax.experimental.pallas.ops.tpu.random.threefry.BLOCK_SIZE", "jax.experimental.pallas.ops.tpu.random.threefry.Shape", "jax.experimental.pallas.ops.tpu.random.threefry.plthreefry_prng_impl", "jax.experimental.pallas.ops.tpu.random.threefry.plthreefry_random_bits", "jax.experimental.pallas.ops.tpu.random.threefry.threefry_2x32_count" ], "n_typable": 11, "n_typed": 6, "n_any": 0, "n_untyped": 5, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 7, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 3, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py", "symbol_reports": [ { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.QKVLayout", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "QKVLayout.HEAD_DIM_MINOR", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "QKVLayout.SEQ_MINOR", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.make_splash_mha_single_device", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.BlockSizes", "methods": [ { "kind": "function", "name": "BlockSizes.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BlockSizes.get_default", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "BlockSizes.has_backward_blocks", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "BlockSizes.block_q", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_kv", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_kv_compute", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_q_dkv", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_kv_dkv", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_kv_dkv_compute", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_q_dq", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.block_kv_dq", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.use_fused_bwd_kernel", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.q_layout", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.k_layout", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSizes.v_layout", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 3, "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 0, "n_attrs": 12, "n_properties": 1 }, { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.SegmentIds", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "SegmentIds.q", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SegmentIds.kv", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.make_splash_mqa", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.make_masked_mha_reference", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.make_splash_mqa_single_device", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.make_masked_mqa_reference", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.make_splash_mha", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.from_head_minor", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.MaskFunctionType", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.NUM_LANES", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.attention_reference_custom", "n_typed": 10, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 11, "n_params": 10, "n_function_overloads": 1, "n_function_params": 10 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.SplashCustomReturnType", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.DEFAULT_MASK_VALUE", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.NUM_SUBLANES", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.SplashResidualsType", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.NN_DIM_NUMBERS", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.make_attention_reference", "n_typed": 4, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.SplashAttentionKernel", "methods": [ { "kind": "function", "name": "SplashAttentionKernel.__init__", "n_typed": 3, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "SplashAttentionKernel.__call__", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "SplashAttentionKernel.manual_sharding_spec", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SplashAttentionKernel.tree_flatten", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "SplashAttentionKernel.tree_unflatten", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "SplashAttentionKernel.kwargs", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "SplashAttentionKernel.fwd_mask_info", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "SplashAttentionKernel.dq_mask_info", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "SplashAttentionKernel.dkv_mask_info", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 18, "n_typed": 5, "n_any": 0, "n_untyped": 13, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 9, "n_attrs": 4, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.NT_DIM_NUMBERS", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.get_kernel_name", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.flash_attention_kernel", "n_typed": 11, "n_any": 0, "n_untyped": 17, "n_overloads": 1, "n_typable": 28, "n_params": 27, "n_function_overloads": 1, "n_function_params": 27 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.attention_reference", "n_typed": 11, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 11, "n_params": 10, "n_function_overloads": 1, "n_function_params": 10 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel", "names": [ "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.BlockSizes", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.DEFAULT_MASK_VALUE", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.MaskFunctionType", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.NN_DIM_NUMBERS", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.NT_DIM_NUMBERS", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.NUM_LANES", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.NUM_SUBLANES", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.QKVLayout", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.SegmentIds", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.SplashAttentionKernel", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.SplashCustomReturnType", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.SplashResidualsType", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.attention_reference", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.attention_reference_custom", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.flash_attention_kernel", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.from_head_minor", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.get_kernel_name", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.make_attention_reference", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.make_masked_mha_reference", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.make_masked_mqa_reference", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.make_splash_mha", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.make_splash_mha_single_device", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.make_splash_mqa", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_kernel.make_splash_mqa_single_device" ], "n_typable": 93, "n_typed": 51, "n_any": 1, "n_untyped": 41, "n_functions": 6, "n_function_overloads": 6, "n_function_params": 57, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 9, "n_classes": 4, "n_attrs": 34, "n_properties": 1, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask.py", "symbol_reports": [ { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.CausalMask", "methods": [ { "kind": "function", "name": "CausalMask.__init__", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "CausalMask.__eq__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "CausalMask.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "CausalMask.offset", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 8, "n_typed": 5, "n_any": 0, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 4, "n_attrs": 1, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.make_local_attention_mask", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.make_causal_mask", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.NumpyMask", "methods": [ { "kind": "function", "name": "NumpyMask.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "NumpyMask.__getitem__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "NumpyMask.__eq__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "NumpyMask.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "NumpyMask.shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "NumpyMask.array", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 7, "n_typed": 3, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 2, "n_attrs": 1, "n_properties": 1 }, { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.LocalMask", "methods": [ { "kind": "function", "name": "LocalMask.__init__", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "LocalMask.__eq__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "LocalMask.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "LocalMask.window_size", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "LocalMask.offset", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 10, "n_typed": 7, "n_any": 0, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 5, "n_attrs": 2, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.FullMask", "methods": [ { "kind": "function", "name": "FullMask.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "FullMask.__getitem__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FullMask.__eq__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "FullMask.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "FullMask.shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "FullMask._shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 7, "n_typed": 3, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 2, "n_attrs": 1, "n_properties": 1 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.make_random_mask", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.MultiHeadMask", "methods": [ { "kind": "function", "name": "MultiHeadMask.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "MultiHeadMask.__getitem__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "MultiHeadMask.__eq__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "MultiHeadMask.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "MultiHeadMask.shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "MultiHeadMask.masks", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 7, "n_typed": 3, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 2, "n_attrs": 1, "n_properties": 1 }, { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.Mask", "methods": [ { "kind": "function", "name": "Mask.__getitem__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Mask.__bool__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Mask.__or__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Mask.__and__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "Mask.shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [], "n_typable": 8, "n_typed": 7, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 3, "n_attrs": 0, "n_properties": 1 }, { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.LogicalAnd", "methods": [ { "kind": "function", "name": "LogicalAnd.__init__", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "LogicalAnd.__getitem__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "LogicalAnd.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "LogicalAnd.shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "LogicalAnd.left", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LogicalAnd.right", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 7, "n_typed": 4, "n_any": 0, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 3, "n_attrs": 2, "n_properties": 1 }, { "kind": "function", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.make_chunk_attention_mask", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.LogicalOr", "methods": [ { "kind": "function", "name": "LogicalOr.__init__", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "LogicalOr.__getitem__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "LogicalOr.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "LogicalOr.shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "LogicalOr.left", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LogicalOr.right", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 7, "n_typed": 4, "n_any": 0, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 3, "n_attrs": 2, "n_properties": 1 }, { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.ChunkedCausalMask", "methods": [ { "kind": "function", "name": "ChunkedCausalMask.__init__", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "ChunkedCausalMask.__eq__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ChunkedCausalMask.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "ChunkedCausalMask.chunk_size", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 8, "n_typed": 5, "n_any": 0, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 4, "n_attrs": 1, "n_properties": 0 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask", "names": [ "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.CausalMask", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.ChunkedCausalMask", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.FullMask", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.LocalMask", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.LogicalAnd", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.LogicalOr", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.Mask", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.MultiHeadMask", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.NumpyMask", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.make_causal_mask", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.make_chunk_attention_mask", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.make_local_attention_mask", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask.make_random_mask" ], "n_typable": 83, "n_typed": 55, "n_any": 0, "n_untyped": 28, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 10, "n_methods": 31, "n_method_overloads": 31, "n_method_params": 28, "n_classes": 9, "n_attrs": 11, "n_properties": 6, "n_type_ignores": 0 }, { "path": "jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask_info.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask_info.process_dynamic_mask_dkv", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "class", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask_info.MaskInfo", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "MaskInfo.data_next", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MaskInfo.mask_next", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MaskInfo.block_mask", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MaskInfo.partial_mask_blocks", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MaskInfo.q_sequence", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MaskInfo.is_dynamic_mask", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 6, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask_info.process_mask_dkv", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask_info.process_mask", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask_info.process_dynamic_mask", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [], "name": "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask_info", "names": [ "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask_info.MaskInfo", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask_info.process_dynamic_mask", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask_info.process_dynamic_mask_dkv", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask_info.process_mask", "jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask_info.process_mask_dkv" ], "n_typable": 4, "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 10, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/mosaic/primitives.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.matmul_pop", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.stochastic_round", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.matmul_acc_lhs", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.unpack_elementwise", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.store", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.async_copy", "n_typed": 3, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.trace_value", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.make_async_copy", "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.touch", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.load", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.pack_elementwise", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.bitcast", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.make_async_remote_copy", "n_typed": 3, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.roll", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.async_remote_copy", "n_typed": 2, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.prng_random_bits", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.matmul_push_rhs", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.prng_seed", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.with_memory_space_constraint", "n_typed": 2, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic.primitives.get_barrier_semaphore", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "type_ignores": [ { "kind": "type", "rules": [ "method-assign" ] }, { "kind": "type", "rules": [ "method-assign" ] } ], "name": "jax._src.pallas.mosaic.primitives", "names": [ "jax._src.pallas.mosaic.primitives.async_copy", "jax._src.pallas.mosaic.primitives.async_remote_copy", "jax._src.pallas.mosaic.primitives.bitcast", "jax._src.pallas.mosaic.primitives.get_barrier_semaphore", "jax._src.pallas.mosaic.primitives.load", "jax._src.pallas.mosaic.primitives.make_async_copy", "jax._src.pallas.mosaic.primitives.make_async_remote_copy", "jax._src.pallas.mosaic.primitives.matmul_acc_lhs", "jax._src.pallas.mosaic.primitives.matmul_pop", "jax._src.pallas.mosaic.primitives.matmul_push_rhs", "jax._src.pallas.mosaic.primitives.pack_elementwise", "jax._src.pallas.mosaic.primitives.prng_random_bits", "jax._src.pallas.mosaic.primitives.prng_seed", "jax._src.pallas.mosaic.primitives.roll", "jax._src.pallas.mosaic.primitives.stochastic_round", "jax._src.pallas.mosaic.primitives.store", "jax._src.pallas.mosaic.primitives.touch", "jax._src.pallas.mosaic.primitives.trace_value", "jax._src.pallas.mosaic.primitives.unpack_elementwise", "jax._src.pallas.mosaic.primitives.with_memory_space_constraint" ], "n_typable": 79, "n_typed": 47, "n_any": 1, "n_untyped": 31, "n_functions": 20, "n_function_overloads": 20, "n_function_params": 59, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 2 }, { "path": "jax/_src/pallas/mosaic/core.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.mosaic.core.create_tensorcore_mesh", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax._src.pallas.mosaic.core.CoreType", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "CoreType.TC", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CoreType.SC_SCALAR_SUBCORE", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CoreType.SC_VECTOR_SUBCORE", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic.core.dma_semaphore", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic.core.GridDimensionSemantics", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "GridDimensionSemantics.PARALLEL", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GridDimensionSemantics.CORE_PARALLEL", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GridDimensionSemantics.SUBCORE_PARALLEL", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GridDimensionSemantics.ARBITRARY", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 4, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic.core.MemorySpace", "methods": [ { "kind": "function", "name": "MemorySpace.__str__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "MemorySpace.from_type", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "MemorySpace.__call__", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "MemorySpace.__getattr__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "MemorySpace.VMEM", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MemorySpace.VMEM_SHARED", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MemorySpace.SMEM", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MemorySpace.CMEM", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MemorySpace.SEMAPHORE", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MemorySpace.HBM", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "MemorySpace.HOST", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 8, "n_typed": 3, "n_any": 0, "n_untyped": 5, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 4, "n_attrs": 7, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic.core.PrefetchScalarGridSpec", "methods": [ { "kind": "function", "name": "PrefetchScalarGridSpec.__init__", "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "PrefetchScalarGridSpec._make_scalar_ref_aval", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "PrefetchScalarGridSpec.num_scalar_prefetch", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 8, "n_typed": 5, "n_any": 0, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 6, "n_attrs": 1, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic.core.CompilerParams", "methods": [ { "kind": "function", "name": "CompilerParams.__init__", "n_typed": 15, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 16, "n_params": 15, "n_function_overloads": 1, "n_function_params": 15 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "CompilerParams.dimension_semantics", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.allow_input_fusion", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.vmem_limit_bytes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.collective_id", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.has_side_effects", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.flags", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.internal_scratch_in_bytes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.serialization_format", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.kernel_type", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.disable_bounds_checks", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.disable_semaphore_checks", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.skip_device_barrier", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.allow_collective_id_without_custom_barrier", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.shape_invariant_numerics", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.use_tc_tiling_on_sc", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.replace", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 16, "n_typed": 15, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 15, "n_attrs": 16, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic.core.SideEffectType", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "SideEffectType.PURE", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SideEffectType.DATAFLOW_SIDE_EFFECTING", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SideEffectType.SIDE_EFFECTING", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic.core.SemaphoreType", "methods": [ { "kind": "function", "name": "SemaphoreType.__call__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SemaphoreType.get_array_aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "SemaphoreType.get_ref_aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "SemaphoreType.REGULAR", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SemaphoreType.DMA", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SemaphoreType.BARRIER", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 4, "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 1, "n_attrs": 3, "n_properties": 0 } ], "type_ignores": [], "name": "jax._src.pallas.mosaic.core", "names": [ "jax._src.pallas.mosaic.core.CompilerParams", "jax._src.pallas.mosaic.core.CoreType", "jax._src.pallas.mosaic.core.GridDimensionSemantics", "jax._src.pallas.mosaic.core.MemorySpace", "jax._src.pallas.mosaic.core.PrefetchScalarGridSpec", "jax._src.pallas.mosaic.core.SemaphoreType", "jax._src.pallas.mosaic.core.SideEffectType", "jax._src.pallas.mosaic.core.create_tensorcore_mesh", "jax._src.pallas.mosaic.core.dma_semaphore" ], "n_typable": 40, "n_typed": 30, "n_any": 0, "n_untyped": 10, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 3, "n_methods": 10, "n_method_overloads": 10, "n_method_params": 26, "n_classes": 8, "n_attrs": 37, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/mosaic/lowering.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.pallas.mosaic.lowering.LoweringException", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 } ], "type_ignores": [], "name": "jax._src.pallas.mosaic.lowering", "names": [ "jax._src.pallas.mosaic.lowering.LoweringException" ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/mosaic/random.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.pallas.mosaic.random.stateful_bernoulli", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.pallas.mosaic.random.stateful_uniform", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.pallas.mosaic.random.sample_block", "n_typed": 7, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "attr", "name": "jax._src.pallas.mosaic.random.stateful_bits", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.pallas.mosaic.random.stateful_normal", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.pallas.mosaic.random.to_pallas_key", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [ { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null } ], "name": "jax._src.pallas.mosaic.random", "names": [ "jax._src.pallas.mosaic.random.sample_block", "jax._src.pallas.mosaic.random.stateful_bernoulli", "jax._src.pallas.mosaic.random.stateful_bits", "jax._src.pallas.mosaic.random.stateful_normal", "jax._src.pallas.mosaic.random.stateful_uniform", "jax._src.pallas.mosaic.random.to_pallas_key" ], "n_typable": 14, "n_typed": 9, "n_any": 0, "n_untyped": 5, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 8, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 4, "n_properties": 0, "n_type_ignores": 4 }, { "path": "jax/_src/pallas/mosaic/helpers.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.mosaic.helpers.sync_copy", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.mosaic.helpers.core_barrier", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic.helpers.run_on_first_core", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax._src.pallas.mosaic.helpers", "names": [ "jax._src.pallas.mosaic.helpers.core_barrier", "jax._src.pallas.mosaic.helpers.run_on_first_core", "jax._src.pallas.mosaic.helpers.sync_copy" ], "n_typable": 9, "n_typed": 4, "n_any": 0, "n_untyped": 5, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 6, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.mosaic.interpret.interpret_pallas_call.reset_tpu_interpret_mode_state", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.pallas.mosaic.interpret.interpret_pallas_call.set_tpu_interpret_mode", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax._src.pallas.mosaic.interpret.interpret_pallas_call.InterpretParams", "methods": [], "properties": [ { "kind": "property", "name": "InterpretParams.num_cores_per_device", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "InterpretParams.dma_execution_mode", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "InterpretParams.random_seed", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "InterpretParams.grid_point_recorder", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "InterpretParams.allow_hbm_allocation_in_run_scoped", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 4, "n_properties": 1 }, { "kind": "function", "name": "jax._src.pallas.mosaic.interpret.interpret_pallas_call.force_tpu_interpret_mode", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax._src.pallas.mosaic.interpret.interpret_pallas_call", "names": [ "jax._src.pallas.mosaic.interpret.interpret_pallas_call.InterpretParams", "jax._src.pallas.mosaic.interpret.interpret_pallas_call.force_tpu_interpret_mode", "jax._src.pallas.mosaic.interpret.interpret_pallas_call.reset_tpu_interpret_mode_state", "jax._src.pallas.mosaic.interpret.interpret_pallas_call.set_tpu_interpret_mode" ], "n_typable": 6, "n_typed": 3, "n_any": 0, "n_untyped": 3, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 2, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 4, "n_properties": 1, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/mosaic/tpu_info.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.mosaic.tpu_info.get_tpu_info", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic.tpu_info.ChipVersion", "methods": [ { "kind": "function", "name": "ChipVersion.__str__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "ChipVersion.TPU_V2", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ChipVersion.TPU_V3", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ChipVersion.TPU_V4I", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ChipVersion.TPU_V4", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ChipVersion.TPU_V5E", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ChipVersion.TPU_V5P", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ChipVersion.TPU_V6E", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ChipVersion.TPU_7X", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 8, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic.tpu_info.TpuInfo", "methods": [ { "kind": "function", "name": "TpuInfo.is_matmul_supported", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "TpuInfo.get_sublane_tiling", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "TpuInfo.is_lite", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "TpuInfo.is_split_chip", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "TpuInfo.chip_version", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TpuInfo.generation", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TpuInfo.num_cores", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TpuInfo.num_lanes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TpuInfo.num_sublanes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TpuInfo.mxu_column_size", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TpuInfo.vmem_capacity_bytes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TpuInfo.cmem_capacity_bytes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TpuInfo.smem_capacity_bytes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TpuInfo.hbm_capacity_bytes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TpuInfo.mem_bw_bytes_per_second", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TpuInfo.bf16_ops_per_second", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TpuInfo.int8_ops_per_second", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TpuInfo.fp8_ops_per_second", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TpuInfo.int4_ops_per_second", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "TpuInfo.sparse_core", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 7, "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 3, "n_attrs": 16, "n_properties": 2 }, { "kind": "class", "name": "jax._src.pallas.mosaic.tpu_info.Tiling", "methods": [], "properties": [ { "kind": "property", "name": "Tiling.shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "Tiling.COMPACT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Tiling.SPARSE_CORE", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 1 }, { "kind": "function", "name": "jax._src.pallas.mosaic.tpu_info.is_tpu_device", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "type_ignores": [], "name": "jax._src.pallas.mosaic.tpu_info", "names": [ "jax._src.pallas.mosaic.tpu_info.ChipVersion", "jax._src.pallas.mosaic.tpu_info.Tiling", "jax._src.pallas.mosaic.tpu_info.TpuInfo", "jax._src.pallas.mosaic.tpu_info.get_tpu_info", "jax._src.pallas.mosaic.tpu_info.is_tpu_device" ], "n_typable": 11, "n_typed": 11, "n_any": 0, "n_untyped": 0, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 3, "n_classes": 3, "n_attrs": 26, "n_properties": 3, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/mosaic/pipeline.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.mosaic.pipeline.emit_pipeline", "n_typed": 8, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 12, "n_params": 11, "n_function_overloads": 1, "n_function_params": 11 }, { "kind": "class", "name": "jax._src.pallas.mosaic.pipeline.BufferedRef", "methods": [ { "kind": "function", "name": "BufferedRef.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BufferedRef.buffer_types", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BufferedRef.create", "n_typed": 4, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "BufferedRef.input", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "BufferedRef.output", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "BufferedRef.accumulator", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "BufferedRef.input_output", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "BufferedRef.with_spec", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BufferedRef.with_next_fetch", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BufferedRef.with_slot_index", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "BufferedRef.bind_existing_ref", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "BufferedRef.unbind_refs", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BufferedRef.compute_slice", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BufferedRef.init_slots", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BufferedRef.advance_copy_in_slot", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BufferedRef.advance_wait_in_slot", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BufferedRef.advance_copy_out_slot", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BufferedRef.advance_wait_out_slot", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BufferedRef.load_slots", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BufferedRef.save_slots", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BufferedRef.copy_in", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "BufferedRef.copy_out", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "BufferedRef.wait_in", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "BufferedRef.wait_out", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "BufferedRef.set_accumulator", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BufferedRef.accumulate", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "BufferedRef.spec", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "BufferedRef.buffer_type", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "BufferedRef.is_buffered", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "BufferedRef.use_lookahead", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "BufferedRef.buffer_count", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "BufferedRef.block_shape", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "BufferedRef.compute_index", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "BufferedRef.current_ref", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "BufferedRef.cumulative_copy_in", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "BufferedRef.current_copy_in_slot", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "BufferedRef.cumulative_copy_out", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "BufferedRef.current_copy_out_slot", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "BufferedRef.cumulative_wait_in", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "BufferedRef.current_wait_in_slot", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "BufferedRef.cumulative_wait_out", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "BufferedRef.current_wait_out_slot", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "BufferedRef.next_fetch_indices", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "BufferedRef._spec", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BufferedRef._buffer_type", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BufferedRef.window_ref", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BufferedRef.accum_ref", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BufferedRef.copy_in_slot", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BufferedRef.wait_in_slot", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BufferedRef.copy_out_slot", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BufferedRef.wait_out_slot", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BufferedRef._copy_in_slot_reg", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BufferedRef._wait_in_slot_reg", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BufferedRef._copy_out_slot_reg", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BufferedRef._wait_out_slot_reg", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BufferedRef.next_fetch_smem", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BufferedRef.next_fetch_sreg", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BufferedRef.sem_recvs", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BufferedRef.sem_sends", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BufferedRef.swap", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BufferedRef.tiling", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 92, "n_typed": 27, "n_any": 0, "n_untyped": 65, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 26, "n_method_overloads": 26, "n_method_params": 49, "n_attrs": 18, "n_properties": 17 }, { "kind": "class", "name": "jax._src.pallas.mosaic.pipeline.BufferedRefBase", "methods": [ { "kind": "function", "name": "BufferedRefBase.init_slots", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BufferedRefBase.advance_copy_in_slot", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BufferedRefBase.advance_wait_in_slot", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BufferedRefBase.advance_copy_out_slot", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BufferedRefBase.advance_wait_out_slot", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BufferedRefBase.load_slots", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BufferedRefBase.save_slots", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BufferedRefBase.get_dma_slice", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "BufferedRefBase.bind_existing_ref", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "BufferedRefBase.unbind_refs", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BufferedRefBase.with_spec", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "BufferedRefBase.spec", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "BufferedRefBase.buffer_type", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "BufferedRefBase.is_buffered", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "BufferedRefBase.is_input", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "BufferedRefBase.is_output", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "BufferedRefBase.is_accumulator", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "BufferedRefBase.is_input_output", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "BufferedRefBase.is_manual", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "BufferedRefBase.block_shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "BufferedRefBase.compute_index", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [], "n_typable": 32, "n_typed": 17, "n_any": 0, "n_untyped": 15, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 11, "n_method_overloads": 11, "n_method_params": 11, "n_attrs": 0, "n_properties": 10 }, { "kind": "function", "name": "jax._src.pallas.mosaic.pipeline.emit_pipeline_with_allocations", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.pallas.mosaic.pipeline.get_pipeline_schedule", "n_typed": 0, "n_any": 1, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.pallas.mosaic.pipeline.make_pipeline_allocations", "n_typed": 1, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 } ], "type_ignores": [], "name": "jax._src.pallas.mosaic.pipeline", "names": [ "jax._src.pallas.mosaic.pipeline.BufferedRef", "jax._src.pallas.mosaic.pipeline.BufferedRefBase", "jax._src.pallas.mosaic.pipeline.emit_pipeline", "jax._src.pallas.mosaic.pipeline.emit_pipeline_with_allocations", "jax._src.pallas.mosaic.pipeline.get_pipeline_schedule", "jax._src.pallas.mosaic.pipeline.make_pipeline_allocations" ], "n_typable": 152, "n_typed": 53, "n_any": 1, "n_untyped": 98, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 24, "n_methods": 37, "n_method_overloads": 37, "n_method_params": 60, "n_classes": 2, "n_attrs": 18, "n_properties": 27, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/einshape.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.einshape.einshape", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [ { "kind": "pyrefly", "rules": [ "bad-override" ] } ], "name": "jax._src.pallas.einshape", "names": [ "jax._src.pallas.einshape.einshape" ], "n_typable": 5, "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 4, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 1 }, { "path": "jax/_src/pallas/mosaic/sc_primitives.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.pack", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.all_reduce_population_count", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.store_compressed", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax._src.pallas.mosaic.sc_primitives.PackFormat", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "PackFormat.COMPRESSED", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "PackFormat.INTERLEAVED", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.unpack", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.subcore_barrier", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.parallel_loop", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 2, "n_typable": 6, "n_params": 5, "n_function_overloads": 2, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.cumsum", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.load_gather", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.fetch_and_add", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.addupdate_compressed", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.store_scatter", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.bitcast", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.addupdate", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.sort_key_val", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.cummax", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.addupdate_scatter", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.all_reduce_ffs", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.load_expanded", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_primitives.scan_count", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [ { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null } ], "name": "jax._src.pallas.mosaic.sc_primitives", "names": [ "jax._src.pallas.mosaic.sc_primitives.PackFormat", "jax._src.pallas.mosaic.sc_primitives.addupdate", "jax._src.pallas.mosaic.sc_primitives.addupdate_compressed", "jax._src.pallas.mosaic.sc_primitives.addupdate_scatter", "jax._src.pallas.mosaic.sc_primitives.all_reduce_ffs", "jax._src.pallas.mosaic.sc_primitives.all_reduce_population_count", "jax._src.pallas.mosaic.sc_primitives.bitcast", "jax._src.pallas.mosaic.sc_primitives.cummax", "jax._src.pallas.mosaic.sc_primitives.cumsum", "jax._src.pallas.mosaic.sc_primitives.fetch_and_add", "jax._src.pallas.mosaic.sc_primitives.load_expanded", "jax._src.pallas.mosaic.sc_primitives.load_gather", "jax._src.pallas.mosaic.sc_primitives.pack", "jax._src.pallas.mosaic.sc_primitives.parallel_loop", "jax._src.pallas.mosaic.sc_primitives.scan_count", "jax._src.pallas.mosaic.sc_primitives.sort_key_val", "jax._src.pallas.mosaic.sc_primitives.store_compressed", "jax._src.pallas.mosaic.sc_primitives.store_scatter", "jax._src.pallas.mosaic.sc_primitives.subcore_barrier", "jax._src.pallas.mosaic.sc_primitives.unpack" ], "n_typable": 71, "n_typed": 70, "n_any": 0, "n_untyped": 1, "n_functions": 19, "n_function_overloads": 20, "n_function_params": 52, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 2, "n_properties": 0, "n_type_ignores": 5 }, { "path": "jax/_src/pallas/mosaic/sc_core.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.pallas.mosaic.sc_core.VectorSubcoreMesh", "methods": [ { "kind": "function", "name": "VectorSubcoreMesh.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "VectorSubcoreMesh.discharges_effect", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "VectorSubcoreMesh.kernel_type", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "VectorSubcoreMesh.default_memory_space", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "VectorSubcoreMesh.shape", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "VectorSubcoreMesh.dimension_semantics", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "VectorSubcoreMesh.core_axis_name", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "VectorSubcoreMesh.subcore_axis_name", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "VectorSubcoreMesh.num_cores", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "VectorSubcoreMesh.num_subcores", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 7, "n_typed": 3, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 1, "n_attrs": 4, "n_properties": 4 }, { "kind": "function", "name": "jax._src.pallas.mosaic.sc_core.get_sparse_core_info", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic.sc_core.MemoryRef", "methods": [ { "kind": "function", "name": "MemoryRef.__init__", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "MemoryRef.get_ref_aval", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "MemoryRef.tiling", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 6, "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 4, "n_attrs": 1, "n_properties": 0 }, { "kind": "class", "name": "jax._src.pallas.mosaic.sc_core.ScalarSubcoreMesh", "methods": [ { "kind": "function", "name": "ScalarSubcoreMesh.discharges_effect", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "ScalarSubcoreMesh.kernel_type", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "ScalarSubcoreMesh.default_memory_space", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "ScalarSubcoreMesh.shape", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "ScalarSubcoreMesh.dimension_semantics", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "ScalarSubcoreMesh.axis_name", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ScalarSubcoreMesh.num_cores", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 6, "n_typed": 3, "n_any": 0, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 1, "n_attrs": 2, "n_properties": 4 }, { "kind": "class", "name": "jax._src.pallas.mosaic.sc_core.BlockSpec", "methods": [ { "kind": "function", "name": "BlockSpec.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BlockSpec.to_block_mapping", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "BlockSpec.indexed_by", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BlockSpec.indexed_dim", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 9, "n_typed": 8, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 7, "n_attrs": 2, "n_properties": 0 } ], "type_ignores": [], "name": "jax._src.pallas.mosaic.sc_core", "names": [ "jax._src.pallas.mosaic.sc_core.BlockSpec", "jax._src.pallas.mosaic.sc_core.MemoryRef", "jax._src.pallas.mosaic.sc_core.ScalarSubcoreMesh", "jax._src.pallas.mosaic.sc_core.VectorSubcoreMesh", "jax._src.pallas.mosaic.sc_core.get_sparse_core_info" ], "n_typable": 29, "n_typed": 20, "n_any": 0, "n_untyped": 9, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 0, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 13, "n_classes": 4, "n_attrs": 9, "n_properties": 8, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/triton/core.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.pallas.triton.core.CompilerParams", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "CompilerParams.num_warps", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "CompilerParams.num_stages", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 } ], "type_ignores": [], "name": "jax._src.pallas.triton.core", "names": [ "jax._src.pallas.triton.core.CompilerParams" ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 2, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/pallas/triton/primitives.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.pallas.triton.primitives.approx_tanh", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.pallas.triton.primitives.debug_barrier", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.pallas.triton.primitives.elementwise_inline_asm", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.pallas.triton.primitives.load", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.pallas.triton.primitives.store", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax._src.pallas.triton.primitives", "names": [ "jax._src.pallas.triton.primitives.approx_tanh", "jax._src.pallas.triton.primitives.debug_barrier", "jax._src.pallas.triton.primitives.elementwise_inline_asm", "jax._src.pallas.triton.primitives.load", "jax._src.pallas.triton.primitives.store" ], "n_typable": 21, "n_typed": 21, "n_any": 0, "n_untyped": 0, "n_functions": 5, "n_function_overloads": 5, "n_function_params": 16, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/profiler.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.profiler.get_profiled_instructions_proto", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.profiler", "names": [ "jax.experimental.profiler.get_profiled_instructions_proto" ], "n_typable": 2, "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 1, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/stateful_rng.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.stateful_rng.stateful_rng", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.stateful_rng.StatefulPRNG", "methods": [ { "kind": "function", "name": "StatefulPRNG.__post_init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "StatefulPRNG.key", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "StatefulPRNG.random", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "StatefulPRNG.uniform", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "StatefulPRNG.normal", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "StatefulPRNG.integers", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "StatefulPRNG.split", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "StatefulPRNG.spawn", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "StatefulPRNG._base_key", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "StatefulPRNG._counter", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 25, "n_typed": 23, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 8, "n_method_overloads": 8, "n_method_params": 17, "n_attrs": 2, "n_properties": 0 } ], "type_ignores": [], "name": "jax._src.stateful_rng", "names": [ "jax._src.stateful_rng.StatefulPRNG", "jax._src.stateful_rng.stateful_rng" ], "n_typable": 28, "n_typed": 26, "n_any": 0, "n_untyped": 2, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 2, "n_methods": 8, "n_method_overloads": 8, "n_method_params": 17, "n_classes": 1, "n_attrs": 2, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/rnn.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.rnn.rnn_fwd_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.rnn.unpack_lstm_weights", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.experimental.rnn.swap_lstm_gates", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.experimental.rnn.rnn_abstract_eval", "n_typed": 6, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 12, "n_params": 11, "n_function_overloads": 1, "n_function_params": 11 }, { "kind": "function", "name": "jax.experimental.rnn.init_lstm_weight", "n_typed": 4, "n_any": 1, "n_untyped": 1, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "attr", "name": "jax.experimental.rnn.rnn_bwd_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.rnn.rnn_bwd_abstract_eval", "n_typed": 6, "n_any": 0, "n_untyped": 11, "n_overloads": 1, "n_typable": 17, "n_params": 16, "n_function_overloads": 1, "n_function_params": 16 }, { "kind": "function", "name": "jax.experimental.rnn.lstm_bwd", "n_typed": 6, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.experimental.rnn.get_num_params_in_lstm", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.rnn.lstm", "n_typed": 12, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 12, "n_params": 11, "n_function_overloads": 1, "n_function_params": 11 }, { "kind": "function", "name": "jax.experimental.rnn.lstm_ref", "n_typed": 14, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 14, "n_params": 13, "n_function_overloads": 1, "n_function_params": 13 }, { "kind": "attr", "name": "jax.experimental.rnn.gpu_rnn", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.rnn.lstm_fwd", "n_typed": 11, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 12, "n_params": 11, "n_function_overloads": 1, "n_function_params": 11 } ], "type_ignores": [ { "kind": "type", "rules": [ "assignment" ] } ], "name": "jax.experimental.rnn", "names": [ "jax.experimental.rnn.get_num_params_in_lstm", "jax.experimental.rnn.gpu_rnn", "jax.experimental.rnn.init_lstm_weight", "jax.experimental.rnn.lstm", "jax.experimental.rnn.lstm_bwd", "jax.experimental.rnn.lstm_fwd", "jax.experimental.rnn.lstm_ref", "jax.experimental.rnn.rnn_abstract_eval", "jax.experimental.rnn.rnn_bwd_abstract_eval", "jax.experimental.rnn.rnn_bwd_p", "jax.experimental.rnn.rnn_fwd_p", "jax.experimental.rnn.swap_lstm_gates", "jax.experimental.rnn.unpack_lstm_weights" ], "n_typable": 101, "n_typed": 70, "n_any": 1, "n_untyped": 30, "n_functions": 10, "n_function_overloads": 10, "n_function_params": 89, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 3, "n_properties": 0, "n_type_ignores": 1 }, { "path": "jax/experimental/roofline/roofline.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.roofline.roofline.roofline_and_grad", "n_typed": 6, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "class", "name": "jax.experimental.roofline.roofline.RooflineRuleContext", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "RooflineRuleContext.name_stack", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RooflineRuleContext.primitive", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RooflineRuleContext.avals_in", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RooflineRuleContext.avals_out", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RooflineRuleContext.jaxpr_eqn_ctx", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RooflineRuleContext.mesh", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RooflineRuleContext.pin_lhs_in_vmem", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RooflineRuleContext.pin_rhs_in_vmem", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 8, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.roofline.roofline.RooflineResult", "methods": [ { "kind": "function", "name": "RooflineResult.zeros", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "RooflineResult.__add__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "RooflineResult.__mul__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "RooflineResult.__rmul__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "RooflineResult.flops", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RooflineResult.unfused_flops", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RooflineResult.ici_bytes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RooflineResult.ici_latency", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RooflineResult.hbm_bytes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RooflineResult.peak_hbm_bytes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RooflineResult.unfused_hbm_bytes", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 7, "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 3, "n_attrs": 7, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.roofline.roofline.RooflineShape", "methods": [ { "kind": "function", "name": "RooflineShape.from_aval", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "RooflineShape.total_bytes", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "RooflineShape.size", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "RooflineShape.bytes", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "RooflineShape.shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RooflineShape.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 6, "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 2, "n_attrs": 2, "n_properties": 2 }, { "kind": "function", "name": "jax.experimental.roofline.roofline.register_standard_roofline", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.roofline.roofline.roofline", "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.experimental.roofline.roofline.register_roofline", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.roofline.roofline.ValidRooflineDtype", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [], "name": "jax.experimental.roofline.roofline", "names": [ "jax.experimental.roofline.roofline.RooflineResult", "jax.experimental.roofline.roofline.RooflineRuleContext", "jax.experimental.roofline.roofline.RooflineShape", "jax.experimental.roofline.roofline.ValidRooflineDtype", "jax.experimental.roofline.roofline.register_roofline", "jax.experimental.roofline.roofline.register_standard_roofline", "jax.experimental.roofline.roofline.roofline", "jax.experimental.roofline.roofline.roofline_and_grad" ], "n_typable": 34, "n_typed": 30, "n_any": 2, "n_untyped": 2, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 17, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 5, "n_classes": 3, "n_attrs": 18, "n_properties": 2, "n_type_ignores": 0 }, { "path": "jax/experimental/scheduling_groups.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.scheduling_groups.unsafe_map", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.scheduling_groups.xla_metadata_call", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.experimental.scheduling_groups.map", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.scheduling_groups.scheduling_group", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.scheduling_groups.attr_get", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.scheduling_groups.zip", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.scheduling_groups.unsafe_zip", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.scheduling_groups.xla_metadata_call_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.scheduling_groups.dce_jaxpr_xla_metadata_rule", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [], "name": "jax.experimental.scheduling_groups", "names": [ "jax.experimental.scheduling_groups.attr_get", "jax.experimental.scheduling_groups.dce_jaxpr_xla_metadata_rule", "jax.experimental.scheduling_groups.map", "jax.experimental.scheduling_groups.scheduling_group", "jax.experimental.scheduling_groups.unsafe_map", "jax.experimental.scheduling_groups.unsafe_zip", "jax.experimental.scheduling_groups.xla_metadata_call", "jax.experimental.scheduling_groups.xla_metadata_call_p", "jax.experimental.scheduling_groups.zip" ], "n_typable": 11, "n_typed": 3, "n_any": 0, "n_untyped": 8, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 6, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 5, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/serialize_executable.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.serialize_executable.deserialize_and_load", "n_typed": 2, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.experimental.serialize_executable.serialize", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.serialize_executable", "names": [ "jax.experimental.serialize_executable.deserialize_and_load", "jax.experimental.serialize_executable.serialize" ], "n_typable": 8, "n_typed": 3, "n_any": 0, "n_untyped": 5, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 6, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/shard_alike.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.shard_alike.shard_alike", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [], "name": "jax._src.shard_alike", "names": [ "jax._src.shard_alike.shard_alike" ], "n_typable": 3, "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 2, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/shard_map.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.shard_map.shard_map", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "type_ignores": [], "name": "jax.experimental.shard_map", "names": [ "jax.experimental.shard_map.shard_map" ], "n_typable": 6, "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 5, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/source_mapper/common.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.source_mapper.common.compile_with_env", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "class", "name": "jax.experimental.source_mapper.common.SourceMapDump", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "SourceMapDump.source_map", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SourceMapDump.generated_code", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SourceMapDump.pass_name", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.source_mapper.common.Pass", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "Pass.name", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Pass.compile_fn", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Pass.generate_dump", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.source_mapper.common.register_pass", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.source_mapper.common.filter_passes", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.source_mapper.common.all_passes", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "class", "name": "jax.experimental.source_mapper.common.GenerateDumpFn", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.source_mapper.common.CompileFn", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.source_mapper.common.flag_env", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.source_mapper.common", "names": [ "jax.experimental.source_mapper.common.CompileFn", "jax.experimental.source_mapper.common.GenerateDumpFn", "jax.experimental.source_mapper.common.Pass", "jax.experimental.source_mapper.common.SourceMapDump", "jax.experimental.source_mapper.common.all_passes", "jax.experimental.source_mapper.common.compile_with_env", "jax.experimental.source_mapper.common.filter_passes", "jax.experimental.source_mapper.common.flag_env", "jax.experimental.source_mapper.common.register_pass" ], "n_typable": 13, "n_typed": 4, "n_any": 0, "n_untyped": 9, "n_functions": 5, "n_function_overloads": 5, "n_function_params": 8, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 4, "n_attrs": 6, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/source_mapper/generate_map.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.source_mapper.generate_map.generate_sourcemaps", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax.experimental.source_mapper.generate_map.SourceMapGeneratorFn", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 } ], "type_ignores": [], "name": "jax.experimental.source_mapper.generate_map", "names": [ "jax.experimental.source_mapper.generate_map.SourceMapGeneratorFn", "jax.experimental.source_mapper.generate_map.generate_sourcemaps" ], "n_typable": 4, "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 3, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/source_mapper/jaxpr.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.source_mapper.jaxpr.canonicalize_filename", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.source_mapper.jaxpr.compile_jaxpr", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.experimental.source_mapper.jaxpr.make_jaxpr_dump", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [ { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": [ "arg-type" ] } ], "name": "jax.experimental.source_mapper.jaxpr", "names": [ "jax.experimental.source_mapper.jaxpr.canonicalize_filename", "jax.experimental.source_mapper.jaxpr.compile_jaxpr", "jax.experimental.source_mapper.jaxpr.make_jaxpr_dump" ], "n_typable": 11, "n_typed": 3, "n_any": 0, "n_untyped": 8, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 8, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 2 }, { "path": "jax/_src/sourcemap.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.sourcemap.SourceMap", "methods": [ { "kind": "function", "name": "SourceMap.from_json", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SourceMap.to_json", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "SourceMap.version", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SourceMap.sources", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SourceMap.sources_content", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SourceMap.names", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SourceMap.mappings", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 3, "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 1, "n_attrs": 5, "n_properties": 0 }, { "kind": "class", "name": "jax._src.sourcemap.MappingsGenerator", "methods": [ { "kind": "function", "name": "MappingsGenerator.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "MappingsGenerator.new_group", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "MappingsGenerator.new_segment", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "MappingsGenerator.mappings", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [], "n_typable": 5, "n_typed": 1, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 1, "n_attrs": 0, "n_properties": 0 } ], "type_ignores": [], "name": "jax._src.sourcemap", "names": [ "jax._src.sourcemap.MappingsGenerator", "jax._src.sourcemap.SourceMap" ], "n_typable": 8, "n_typed": 4, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 2, "n_classes": 2, "n_attrs": 5, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/source_mapper/mlir.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.source_mapper.mlir.create_mlir_sourcemap", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.source_mapper.mlir.LOC_REGEX", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.source_mapper.mlir.SCOPED_REGEX", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.source_mapper.mlir.Redirect", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.source_mapper.mlir.SRC_REGEX", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.source_mapper.mlir.Location", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.source_mapper.mlir.CALLSITE_REGEX", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.source_mapper.mlir.parse_mlir_locations", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.experimental.source_mapper.mlir", "names": [ "jax.experimental.source_mapper.mlir.CALLSITE_REGEX", "jax.experimental.source_mapper.mlir.LOC_REGEX", "jax.experimental.source_mapper.mlir.Location", "jax.experimental.source_mapper.mlir.Redirect", "jax.experimental.source_mapper.mlir.SCOPED_REGEX", "jax.experimental.source_mapper.mlir.SRC_REGEX", "jax.experimental.source_mapper.mlir.create_mlir_sourcemap", "jax.experimental.source_mapper.mlir.parse_mlir_locations" ], "n_typable": 8, "n_typed": 4, "n_any": 0, "n_untyped": 4, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 2, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 6, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/source_mapper/hlo.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.source_mapper.hlo.stable_hlo_generate_dump", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.source_mapper.hlo.original_hlo_generate_dump", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.source_mapper.hlo.optimized_generate_dump", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax.experimental.source_mapper.hlo.HloPass", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "HloPass.STABLE_HLO", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "HloPass.ORIGINAL", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "HloPass.OPTIMIZED", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.source_mapper.hlo.parse_hlo_dump", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.source_mapper.hlo.METADATA_REGEX", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.source_mapper.hlo.trace_and_lower", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "type_ignores": [ { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": [ "arg-type" ] } ], "name": "jax.experimental.source_mapper.hlo", "names": [ "jax.experimental.source_mapper.hlo.HloPass", "jax.experimental.source_mapper.hlo.METADATA_REGEX", "jax.experimental.source_mapper.hlo.optimized_generate_dump", "jax.experimental.source_mapper.hlo.original_hlo_generate_dump", "jax.experimental.source_mapper.hlo.parse_hlo_dump", "jax.experimental.source_mapper.hlo.stable_hlo_generate_dump", "jax.experimental.source_mapper.hlo.trace_and_lower" ], "n_typable": 19, "n_typed": 9, "n_any": 0, "n_untyped": 10, "n_functions": 5, "n_function_overloads": 5, "n_function_params": 13, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 4, "n_properties": 0, "n_type_ignores": 6 }, { "path": "jax/experimental/sparse/csr.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.sparse.csr.csr_matmat_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.sparse.csr.csr_matvec_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.sparse.csr.csr_matvec", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax.experimental.sparse.csr.csr_fromdense_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.sparse.csr.csr_matmat", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax.experimental.sparse.csr.csr_todense_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "class", "name": "jax.experimental.sparse.csr.CSR", "methods": [ { "kind": "function", "name": "CSR.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "CSR.fromdense", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "CSR._empty", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "CSR._eye", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "CSR.todense", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "CSR.transpose", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "CSR.__matmul__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "CSR.tree_flatten", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "CSR.tree_unflatten", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "CSR.data", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "CSR.indices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "CSR.indptr", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "CSR.shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "CSR.nse", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "CSR.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "CSR._bufs", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 33, "n_typed": 4, "n_any": 0, "n_untyped": 29, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 9, "n_method_overloads": 9, "n_method_params": 17, "n_attrs": 7, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.sparse.csr.csr_fromdense", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax.experimental.sparse.csr.CSC", "methods": [ { "kind": "function", "name": "CSC.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "CSC.fromdense", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "CSC._empty", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "CSC._eye", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "CSC.todense", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "CSC.transpose", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "CSC.__matmul__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "CSC.tree_flatten", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "CSC.tree_unflatten", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "CSC.data", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "CSC.indices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "CSC.indptr", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "CSC.shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "CSC.nse", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "CSC.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 32, "n_typed": 4, "n_any": 0, "n_untyped": 28, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 9, "n_method_overloads": 9, "n_method_params": 17, "n_attrs": 6, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.sparse.csr.csr_todense", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.sparse.csr.Shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [], "name": "jax.experimental.sparse.csr", "names": [ "jax.experimental.sparse.csr.CSC", "jax.experimental.sparse.csr.CSR", "jax.experimental.sparse.csr.Shape", "jax.experimental.sparse.csr.csr_fromdense", "jax.experimental.sparse.csr.csr_fromdense_p", "jax.experimental.sparse.csr.csr_matmat", "jax.experimental.sparse.csr.csr_matmat_p", "jax.experimental.sparse.csr.csr_matvec", "jax.experimental.sparse.csr.csr_matvec_p", "jax.experimental.sparse.csr.csr_todense", "jax.experimental.sparse.csr.csr_todense_p" ], "n_typable": 83, "n_typed": 22, "n_any": 0, "n_untyped": 61, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 10, "n_methods": 18, "n_method_overloads": 18, "n_method_params": 34, "n_classes": 2, "n_attrs": 18, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/sparse/bcsr.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.sparse.bcsr.bcsr_extract_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.sparse.bcsr.bcsr_dot_general", "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "attr", "name": "jax.experimental.sparse.bcsr.bcsr_todense_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.sparse.bcsr.bcsr_sum_duplicates", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.sparse.bcsr.bcsr_broadcast_in_dim", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "attr", "name": "jax.experimental.sparse.bcsr.bcsr_dot_general_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "class", "name": "jax.experimental.sparse.bcsr.BCSR", "methods": [ { "kind": "function", "name": "BCSR.__init__", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "BCSR.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BCSR.transpose", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "BCSR.tree_flatten", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BCSR.tree_unflatten", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "BCSR._empty", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "BCSR.sum_duplicates", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "BCSR.fromdense", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "BCSR.todense", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BCSR.to_bcoo", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BCSR.from_bcoo", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BCSR.from_scipy_sparse", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "properties": [ { "kind": "property", "name": "BCSR._sparse_shape", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "BCSR.data", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "BCSR.indices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "BCSR.indptr", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "BCSR.shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "BCSR.nse", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "BCSR.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "BCSR.n_batch", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "BCSR.n_sparse", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "BCSR.n_dense", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "BCSR.indices_sorted", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "BCSR.unique_indices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "BCSR._bufs", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "BCSR._info", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 52, "n_typed": 16, "n_any": 0, "n_untyped": 36, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 12, "n_method_overloads": 12, "n_method_params": 26, "n_attrs": 13, "n_properties": 1 }, { "kind": "function", "name": "jax.experimental.sparse.bcsr.bcsr_concatenate", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.experimental.sparse.bcsr.bcsr_fromdense_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.sparse.bcsr.bcsr_extract", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.sparse.bcsr.bcsr_todense", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.sparse.bcsr.bcsr_fromdense", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.experimental.sparse.bcsr.bcsr_eliminate_zeros", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax.experimental.sparse.bcsr.BCSRProperties", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "BCSRProperties.n_batch", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BCSRProperties.n_dense", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BCSRProperties.nse", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 } ], "type_ignores": [], "name": "jax.experimental.sparse.bcsr", "names": [ "jax.experimental.sparse.bcsr.BCSR", "jax.experimental.sparse.bcsr.BCSRProperties", "jax.experimental.sparse.bcsr.bcsr_broadcast_in_dim", "jax.experimental.sparse.bcsr.bcsr_concatenate", "jax.experimental.sparse.bcsr.bcsr_dot_general", "jax.experimental.sparse.bcsr.bcsr_dot_general_p", "jax.experimental.sparse.bcsr.bcsr_eliminate_zeros", "jax.experimental.sparse.bcsr.bcsr_extract", "jax.experimental.sparse.bcsr.bcsr_extract_p", "jax.experimental.sparse.bcsr.bcsr_fromdense", "jax.experimental.sparse.bcsr.bcsr_fromdense_p", "jax.experimental.sparse.bcsr.bcsr_sum_duplicates", "jax.experimental.sparse.bcsr.bcsr_todense", "jax.experimental.sparse.bcsr.bcsr_todense_p" ], "n_typable": 89, "n_typed": 47, "n_any": 0, "n_untyped": 42, "n_functions": 8, "n_function_overloads": 8, "n_function_params": 25, "n_methods": 12, "n_method_overloads": 12, "n_method_params": 26, "n_classes": 2, "n_attrs": 20, "n_properties": 1, "n_type_ignores": 0 }, { "path": "jax/experimental/sparse/bcoo.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_rev", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.experimental.sparse.bcoo.bcoo_sort_indices_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_broadcast_in_dim", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_dot_general", "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_slice", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_reshape", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_concatenate", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_extract", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax.experimental.sparse.bcoo.bcoo_spdot_general_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_gather", "n_typed": 8, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_dynamic_slice", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax.experimental.sparse.bcoo.bcoo_fromdense_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.sparse.bcoo.bcoo_extract_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.sparse.bcoo.bcoo_todense_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_squeeze", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_update_layout", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_dot_general_sampled", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_multiply_dense", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_sort_indices", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.sparse.bcoo.bcoo_dot_general_sampled_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_multiply_sparse", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_transpose", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_sum_duplicates", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_conv_general_dilated", "n_typed": 1, "n_any": 0, "n_untyped": 12, "n_overloads": 1, "n_typable": 13, "n_params": 12, "n_function_overloads": 1, "n_function_params": 12 }, { "kind": "attr", "name": "jax.experimental.sparse.bcoo.bcoo_transpose_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "class", "name": "jax.experimental.sparse.bcoo.BCOO", "methods": [ { "kind": "function", "name": "BCOO.__init__", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "BCOO.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BCOO.reshape", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "BCOO.astype", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "BCOO.sum", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BCOO.fromdense", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "BCOO.from_scipy_sparse", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "BCOO._empty", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "BCOO._eye", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "BCOO.update_layout", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "BCOO.sum_duplicates", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "BCOO.sort_indices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BCOO.todense", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BCOO.transpose", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "BCOO.tree_flatten", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "BCOO.tree_unflatten", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "BCOO.data", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "BCOO.indices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "BCOO.shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "BCOO.nse", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "BCOO.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "BCOO.n_batch", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "BCOO.n_sparse", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "BCOO.n_dense", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "BCOO.indices_sorted", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "BCOO.unique_indices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "BCOO._info", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "BCOO._bufs", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 66, "n_typed": 48, "n_any": 0, "n_untyped": 18, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 16, "n_method_overloads": 16, "n_method_params": 38, "n_attrs": 12, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_fromdense", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "attr", "name": "jax.experimental.sparse.bcoo.bcoo_dot_general_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_reduce_sum", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_todense", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.sparse.bcoo.bcoo_sum_duplicates_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.sparse.bcoo.CUSPARSE_DATA_DTYPES", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.experimental.sparse.bcoo.CUSPARSE_INDEX_DTYPES", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.sparse.bcoo.bcoo_eliminate_zeros", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax.experimental.sparse.bcoo.BCOOProperties", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "BCOOProperties.n_batch", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BCOOProperties.n_sparse", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BCOOProperties.n_dense", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BCOOProperties.nse", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 4, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.sparse.bcoo.Buffer", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 } ], "type_ignores": [], "name": "jax.experimental.sparse.bcoo", "names": [ "jax.experimental.sparse.bcoo.BCOO", "jax.experimental.sparse.bcoo.BCOOProperties", "jax.experimental.sparse.bcoo.Buffer", "jax.experimental.sparse.bcoo.CUSPARSE_DATA_DTYPES", "jax.experimental.sparse.bcoo.CUSPARSE_INDEX_DTYPES", "jax.experimental.sparse.bcoo.bcoo_broadcast_in_dim", "jax.experimental.sparse.bcoo.bcoo_concatenate", "jax.experimental.sparse.bcoo.bcoo_conv_general_dilated", "jax.experimental.sparse.bcoo.bcoo_dot_general", "jax.experimental.sparse.bcoo.bcoo_dot_general_p", "jax.experimental.sparse.bcoo.bcoo_dot_general_sampled", "jax.experimental.sparse.bcoo.bcoo_dot_general_sampled_p", "jax.experimental.sparse.bcoo.bcoo_dynamic_slice", "jax.experimental.sparse.bcoo.bcoo_eliminate_zeros", "jax.experimental.sparse.bcoo.bcoo_extract", "jax.experimental.sparse.bcoo.bcoo_extract_p", "jax.experimental.sparse.bcoo.bcoo_fromdense", "jax.experimental.sparse.bcoo.bcoo_fromdense_p", "jax.experimental.sparse.bcoo.bcoo_gather", "jax.experimental.sparse.bcoo.bcoo_multiply_dense", "jax.experimental.sparse.bcoo.bcoo_multiply_sparse", "jax.experimental.sparse.bcoo.bcoo_reduce_sum", "jax.experimental.sparse.bcoo.bcoo_reshape", "jax.experimental.sparse.bcoo.bcoo_rev", "jax.experimental.sparse.bcoo.bcoo_slice", "jax.experimental.sparse.bcoo.bcoo_sort_indices", "jax.experimental.sparse.bcoo.bcoo_sort_indices_p", "jax.experimental.sparse.bcoo.bcoo_spdot_general_p", "jax.experimental.sparse.bcoo.bcoo_squeeze", "jax.experimental.sparse.bcoo.bcoo_sum_duplicates", "jax.experimental.sparse.bcoo.bcoo_sum_duplicates_p", "jax.experimental.sparse.bcoo.bcoo_todense", "jax.experimental.sparse.bcoo.bcoo_todense_p", "jax.experimental.sparse.bcoo.bcoo_transpose", "jax.experimental.sparse.bcoo.bcoo_transpose_p", "jax.experimental.sparse.bcoo.bcoo_update_layout" ], "n_typable": 174, "n_typed": 128, "n_any": 0, "n_untyped": 46, "n_functions": 22, "n_function_overloads": 22, "n_function_params": 77, "n_methods": 16, "n_method_overloads": 16, "n_method_params": 38, "n_classes": 3, "n_attrs": 27, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/sparse/ad.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.sparse.ad.jacrev", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.sparse.ad.jacfwd", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.sparse.ad.value_and_grad", "n_typed": 3, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.sparse.ad.grad", "n_typed": 3, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.sparse.ad.flatten_fun_for_sparse_ad", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax.experimental.sparse.ad.is_sparse", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [], "name": "jax.experimental.sparse.ad", "names": [ "jax.experimental.sparse.ad.flatten_fun_for_sparse_ad", "jax.experimental.sparse.ad.grad", "jax.experimental.sparse.ad.is_sparse", "jax.experimental.sparse.ad.jacfwd", "jax.experimental.sparse.ad.jacrev", "jax.experimental.sparse.ad.value_and_grad" ], "n_typable": 25, "n_typed": 16, "n_any": 0, "n_untyped": 9, "n_functions": 5, "n_function_overloads": 5, "n_function_params": 19, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/sparse/api.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.sparse.api.eye", "n_typed": 7, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax.experimental.sparse.api.todense", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.sparse.api.todense_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.sparse.api.empty", "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "type_ignores": [], "name": "jax.experimental.sparse.api", "names": [ "jax.experimental.sparse.api.empty", "jax.experimental.sparse.api.eye", "jax.experimental.sparse.api.todense", "jax.experimental.sparse.api.todense_p" ], "n_typable": 17, "n_typed": 14, "n_any": 0, "n_untyped": 3, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 13, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/sparse/coo.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.sparse.coo.coo_fromdense_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.sparse.coo.coo_matvec_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.sparse.coo.coo_matmat", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax.experimental.sparse.coo.coo_todense_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.sparse.coo.coo_matmat_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.sparse.coo.coo_matvec", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax.experimental.sparse.coo.COO", "methods": [ { "kind": "function", "name": "COO.__init__", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "COO.fromdense", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "COO._sort_indices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "COO._empty", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "COO._eye", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "COO.todense", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "COO.transpose", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "COO.tree_flatten", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "COO.tree_unflatten", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "COO.__matmul__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "COO.data", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "COO.row", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "COO.col", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "COO.shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "COO.nse", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "COO.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "COO._info", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "COO._bufs", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "COO._rows_sorted", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "COO._cols_sorted", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 39, "n_typed": 31, "n_any": 0, "n_untyped": 8, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 10, "n_method_overloads": 10, "n_method_params": 19, "n_attrs": 10, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.sparse.coo.coo_todense", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.sparse.coo.coo_fromdense", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax.experimental.sparse.coo.COOInfo", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "COOInfo.shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "COOInfo.rows_sorted", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "COOInfo.cols_sorted", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.sparse.coo.Shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [], "name": "jax.experimental.sparse.coo", "names": [ "jax.experimental.sparse.coo.COO", "jax.experimental.sparse.coo.COOInfo", "jax.experimental.sparse.coo.Shape", "jax.experimental.sparse.coo.coo_fromdense", "jax.experimental.sparse.coo.coo_fromdense_p", "jax.experimental.sparse.coo.coo_matmat", "jax.experimental.sparse.coo.coo_matmat_p", "jax.experimental.sparse.coo.coo_matvec", "jax.experimental.sparse.coo.coo_matvec_p", "jax.experimental.sparse.coo.coo_todense", "jax.experimental.sparse.coo.coo_todense_p" ], "n_typable": 57, "n_typed": 45, "n_any": 0, "n_untyped": 12, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 10, "n_methods": 10, "n_method_overloads": 10, "n_method_params": 19, "n_classes": 2, "n_attrs": 18, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/sparse/random.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.sparse.random.random_bcoo", "n_typed": 0, "n_any": 0, "n_untyped": 12, "n_overloads": 1, "n_typable": 12, "n_params": 11, "n_function_overloads": 1, "n_function_params": 11 } ], "type_ignores": [], "name": "jax.experimental.sparse.random", "names": [ "jax.experimental.sparse.random.random_bcoo" ], "n_typable": 12, "n_typed": 0, "n_any": 0, "n_untyped": 12, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 11, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/sparse/util.py", "symbol_reports": [ { "kind": "class", "name": "jax.experimental.sparse.util.CuSparseEfficiencyWarning", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.sparse.util.SparseEfficiencyWarning", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax.experimental.sparse.util.SparseEfficiencyError", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.sparse.util.broadcasting_vmap", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.experimental.sparse.util.nfold_vmap", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "class", "name": "jax.experimental.sparse.util.SparseInfo", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "SparseInfo.shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SparseInfo.indices_sorted", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SparseInfo.unique_indices", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "attr", "name": "jax.experimental.sparse.util.Shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [], "name": "jax.experimental.sparse.util", "names": [ "jax.experimental.sparse.util.CuSparseEfficiencyWarning", "jax.experimental.sparse.util.Shape", "jax.experimental.sparse.util.SparseEfficiencyError", "jax.experimental.sparse.util.SparseEfficiencyWarning", "jax.experimental.sparse.util.SparseInfo", "jax.experimental.sparse.util.broadcasting_vmap", "jax.experimental.sparse.util.nfold_vmap" ], "n_typable": 9, "n_typed": 0, "n_any": 0, "n_untyped": 9, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 7, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 4, "n_attrs": 4, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/sparse/transform.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.sparse.transform.sparsify", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax.experimental.sparse.transform.SparseTracer", "methods": [ { "kind": "function", "name": "SparseTracer.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "SparseTracer.full_lower", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "SparseTracer.spenv", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "SparseTracer.aval", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [], "n_typable": 6, "n_typed": 1, "n_any": 0, "n_untyped": 5, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 2, "n_attrs": 0, "n_properties": 2 }, { "kind": "attr", "name": "jax.experimental.sparse.transform.sparse_rules_bcoo", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.sparse.transform.sparsify_fun", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.experimental.sparse.transform.sparse_rules_bcsr", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "class", "name": "jax.experimental.sparse.transform.SparseTrace", "methods": [ { "kind": "function", "name": "SparseTrace.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "SparseTrace.to_sparse_tracer", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SparseTrace.process_primitive", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "SparseTrace.process_call", "n_typed": 1, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "SparseTrace.process_custom_jvp_call", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "SparseTrace.parent_trace", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "SparseTrace.tag", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "SparseTrace.spenv", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 24, "n_typed": 1, "n_any": 0, "n_untyped": 23, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 16, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.sparse.transform.eval_sparse", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.experimental.sparse.transform.sparsify_subtrace", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.experimental.sparse.transform.spvalues_to_avals", "n_typed": 1, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.sparse.transform.spvalues_to_arrays", "n_typed": 1, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax.experimental.sparse.transform.SparsifyValue", "methods": [ { "kind": "function", "name": "SparsifyValue.is_sparse", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "SparsifyValue.is_dense", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "SparsifyValue.is_bcoo", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "SparsifyValue.is_bcsr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "SparsifyValue.ndim", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "SparsifyValue.shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SparsifyValue.data_ref", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SparsifyValue.indices_ref", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SparsifyValue.indptr_ref", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SparsifyValue.indices_sorted", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SparsifyValue.unique_indices", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 5, "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 0, "n_attrs": 6, "n_properties": 1 }, { "kind": "function", "name": "jax.experimental.sparse.transform.arrays_to_spvalues", "n_typed": 1, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.sparse.transform.sparsify_raw", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax.experimental.sparse.transform.SparsifyEnv", "methods": [ { "kind": "function", "name": "SparsifyEnv.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SparsifyEnv._push", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SparsifyEnv.data", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SparsifyEnv.indices", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SparsifyEnv.indptr", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SparsifyEnv.dense", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SparsifyEnv.sparse", "n_typed": 0, "n_any": 0, "n_untyped": 10, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "SparsifyEnv._buffers", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 23, "n_typed": 5, "n_any": 4, "n_untyped": 14, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 15, "n_attrs": 1, "n_properties": 0 } ], "type_ignores": [], "name": "jax.experimental.sparse.transform", "names": [ "jax.experimental.sparse.transform.SparseTrace", "jax.experimental.sparse.transform.SparseTracer", "jax.experimental.sparse.transform.SparsifyEnv", "jax.experimental.sparse.transform.SparsifyValue", "jax.experimental.sparse.transform.arrays_to_spvalues", "jax.experimental.sparse.transform.eval_sparse", "jax.experimental.sparse.transform.sparse_rules_bcoo", "jax.experimental.sparse.transform.sparse_rules_bcsr", "jax.experimental.sparse.transform.sparsify", "jax.experimental.sparse.transform.sparsify_fun", "jax.experimental.sparse.transform.sparsify_raw", "jax.experimental.sparse.transform.sparsify_subtrace", "jax.experimental.sparse.transform.spvalues_to_arrays", "jax.experimental.sparse.transform.spvalues_to_avals" ], "n_typable": 89, "n_typed": 18, "n_any": 10, "n_untyped": 61, "n_functions": 8, "n_function_overloads": 8, "n_function_params": 21, "n_methods": 18, "n_method_overloads": 18, "n_method_params": 33, "n_classes": 4, "n_attrs": 12, "n_properties": 3, "n_type_ignores": 0 }, { "path": "jax/experimental/sparse/__init__.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.sparse.jacobian", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [], "name": "jax.experimental.sparse", "names": [ "jax.experimental.sparse.jacobian" ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/sparse/_base.py", "symbol_reports": [ { "kind": "class", "name": "jax.experimental.sparse._base.JAXSparse", "methods": [ { "kind": "function", "name": "JAXSparse.__len__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "JAXSparse.__init__", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "JAXSparse.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "JAXSparse.tree_flatten", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "JAXSparse.tree_unflatten", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "JAXSparse.transpose", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "JAXSparse.block_until_ready", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "JAXSparse.sum", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "JAXSparse.__neg__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "JAXSparse.__pos__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "JAXSparse.__matmul__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "JAXSparse.__rmatmul__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "JAXSparse.__mul__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "JAXSparse.__rmul__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "JAXSparse.__add__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "JAXSparse.__radd__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "JAXSparse.__sub__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "JAXSparse.__rsub__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "JAXSparse.__getitem__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "JAXSparse.size", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "JAXSparse.ndim", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "JAXSparse.T", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "JAXSparse.data", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "JAXSparse.shape", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "JAXSparse.nse", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "JAXSparse.dtype", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "JAXSparse.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 42, "n_typed": 8, "n_any": 0, "n_untyped": 34, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 19, "n_method_overloads": 19, "n_method_params": 16, "n_attrs": 5, "n_properties": 3 } ], "type_ignores": [ { "kind": "type", "rules": null } ], "name": "jax.experimental.sparse._base", "names": [ "jax.experimental.sparse._base.JAXSparse" ], "n_typable": 42, "n_typed": 8, "n_any": 0, "n_untyped": 34, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 19, "n_method_overloads": 19, "n_method_params": 16, "n_classes": 1, "n_attrs": 5, "n_properties": 3, "n_type_ignores": 1 }, { "path": "jax/experimental/sparse/linalg.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.sparse.linalg.spsolve", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "attr", "name": "jax.experimental.sparse.linalg.spsolve_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.sparse.linalg.lobpcg_standard", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax.experimental.sparse.linalg", "names": [ "jax.experimental.sparse.linalg.lobpcg_standard", "jax.experimental.sparse.linalg.spsolve", "jax.experimental.sparse.linalg.spsolve_p" ], "n_typable": 13, "n_typed": 4, "n_any": 0, "n_untyped": 9, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 10, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/sparse/nm.py", "symbol_reports": [ { "kind": "attr", "name": "jax.experimental.sparse.nm.nm_pack_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.experimental.sparse.nm.nm_spmm_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax.experimental.sparse.nm.nm_spmm", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.experimental.sparse.nm.nm_pack", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax.experimental.sparse.nm", "names": [ "jax.experimental.sparse.nm.nm_pack", "jax.experimental.sparse.nm.nm_pack_p", "jax.experimental.sparse.nm.nm_spmm", "jax.experimental.sparse.nm.nm_spmm_p" ], "n_typable": 13, "n_typed": 9, "n_any": 0, "n_untyped": 4, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 9, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 2, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/sparse/test_util.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.sparse.test_util.iter_bcsr_layouts", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax.experimental.sparse.test_util.SparseLayout", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "SparseLayout.n_batch", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SparseLayout.n_dense", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SparseLayout.n_sparse", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.sparse.test_util.is_sparse", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.experimental.sparse.test_util.MATMUL_TOL", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.experimental.sparse.test_util.iter_sparse_layouts", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.experimental.sparse.test_util.iter_subsets", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.sparse.test_util.rand_bcoo", "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "class", "name": "jax.experimental.sparse.test_util.BatchedDotGeneralProperties", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "BatchedDotGeneralProperties.lhs_shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BatchedDotGeneralProperties.rhs_shape", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BatchedDotGeneralProperties.n_batch", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BatchedDotGeneralProperties.n_dense", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "BatchedDotGeneralProperties.dimension_numbers", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 5, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.sparse.test_util.rand_sparse", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "class", "name": "jax.experimental.sparse.test_util.SparseTestCase", "methods": [ { "kind": "function", "name": "SparseTestCase.assertSparseArraysEquivalent", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "SparseTestCase._CheckAgainstDense", "n_typed": 0, "n_any": 0, "n_untyped": 10, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "SparseTestCase._CheckGradsSparse", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "SparseTestCase._random_bdims", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "SparseTestCase._CheckBatchingSparse", "n_typed": 0, "n_any": 0, "n_untyped": 12, "n_overloads": 1, "n_typable": 12, "n_params": 11, "n_function_overloads": 1, "n_function_params": 11 } ], "properties": [], "attrs": [], "n_typable": 40, "n_typed": 0, "n_any": 0, "n_untyped": 40, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 35, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.sparse.test_util.rand_bcsr", "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "type_ignores": [], "name": "jax.experimental.sparse.test_util", "names": [ "jax.experimental.sparse.test_util.BatchedDotGeneralProperties", "jax.experimental.sparse.test_util.MATMUL_TOL", "jax.experimental.sparse.test_util.SparseLayout", "jax.experimental.sparse.test_util.SparseTestCase", "jax.experimental.sparse.test_util.is_sparse", "jax.experimental.sparse.test_util.iter_bcsr_layouts", "jax.experimental.sparse.test_util.iter_sparse_layouts", "jax.experimental.sparse.test_util.iter_subsets", "jax.experimental.sparse.test_util.rand_bcoo", "jax.experimental.sparse.test_util.rand_bcsr", "jax.experimental.sparse.test_util.rand_sparse" ], "n_typable": 67, "n_typed": 16, "n_any": 0, "n_untyped": 51, "n_functions": 7, "n_function_overloads": 7, "n_function_params": 20, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 35, "n_classes": 3, "n_attrs": 9, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/topologies.py", "symbol_reports": [ { "kind": "function", "name": "jax.experimental.topologies.get_attached_topology", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.experimental.topologies.make_mesh", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "class", "name": "jax.experimental.topologies.TopologyDescription", "methods": [ { "kind": "function", "name": "TopologyDescription.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "TopologyDescription.devices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 3, "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 1, "n_attrs": 1, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.topologies.get_topology_desc", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax.experimental.topologies", "names": [ "jax.experimental.topologies.TopologyDescription", "jax.experimental.topologies.get_attached_topology", "jax.experimental.topologies.get_topology_desc", "jax.experimental.topologies.make_mesh" ], "n_typable": 14, "n_typed": 11, "n_any": 0, "n_untyped": 3, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 8, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 1, "n_classes": 1, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/experimental/transfer.py", "symbol_reports": [ { "kind": "class", "name": "jax.experimental.transfer.TransferConnection", "methods": [ { "kind": "function", "name": "TransferConnection._pull_flat", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "TransferConnection.pull", "n_typed": 1, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [], "n_typable": 7, "n_typed": 1, "n_any": 2, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 5, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax.experimental.transfer.make_error_array", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax.experimental.transfer.TransferServer", "methods": [ { "kind": "function", "name": "TransferServer.address", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "TransferServer._await_pull_flat", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "TransferServer.connect", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TransferServer.await_pull", "n_typed": 1, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [], "n_typable": 9, "n_typed": 5, "n_any": 2, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 5, "n_attrs": 0, "n_properties": 0 } ], "type_ignores": [], "name": "jax.experimental.transfer", "names": [ "jax.experimental.transfer.TransferConnection", "jax.experimental.transfer.TransferServer", "jax.experimental.transfer.make_error_array" ], "n_typable": 19, "n_typed": 6, "n_any": 4, "n_untyped": 9, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 2, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 10, "n_classes": 2, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/xla_metadata.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.xla_metadata.set_xla_metadata", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [], "name": "jax._src.xla_metadata", "names": [ "jax._src.xla_metadata.set_xla_metadata" ], "n_typable": 3, "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 2, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/extend/backend.py", "symbol_reports": [ { "kind": "attr", "name": "jax.extend.backend.ifrt_proxy", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [ { "kind": "type", "rules": null } ], "name": "jax.extend.backend", "names": [ "jax.extend.backend.ifrt_proxy" ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 1 }, { "path": "jax/_src/interpreters/pxla.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.interpreters.pxla.get_default_device", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.interpreters.pxla.clear_in_memory_compilation_cache", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "attr", "name": "jax._src.interpreters.pxla.xla_pmap_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [ { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null } ], "name": "jax._src.interpreters.pxla", "names": [ "jax._src.interpreters.pxla.clear_in_memory_compilation_cache", "jax._src.interpreters.pxla.get_default_device", "jax._src.interpreters.pxla.xla_pmap_p" ], "n_typable": 3, "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 3 }, { "path": "jax/_src/compiler.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.compiler.get_compile_options", "n_typed": 11, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 12, "n_params": 11, "n_function_overloads": 1, "n_function_params": 11 } ], "type_ignores": [ { "kind": "type", "rules": [ "missing-attribute" ] }, { "kind": "type", "rules": [ "missing-attribute" ] } ], "name": "jax._src.compiler", "names": [ "jax._src.compiler.get_compile_options" ], "n_typable": 12, "n_typed": 11, "n_any": 0, "n_untyped": 1, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 11, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 2 }, { "path": "jax/_src/abstract_arrays.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.abstract_arrays.array_types", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "type_ignores": [], "name": "jax._src.abstract_arrays", "names": [ "jax._src.abstract_arrays.array_types" ], "n_typable": 1, "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/lax/lax.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.lax.lax.rng_uniform_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.eq_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.reduce_min_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.create_token_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.round_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.sub_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.pad_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.concatenate_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.sinh_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.real_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.population_count_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.bitcast_convert_type_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.reduce_max_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.xor_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.empty2_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.rng_bit_generator_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.reshape_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.sqrt_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.mul_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.atanh_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.reduce_prod_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.cosh_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.reduce_xor_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.cbrt_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.complex_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.cos_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.or_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.lt_to_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.conj_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.tanh_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.after_all_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.expm1_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.ge_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.min_p", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.log_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.neg_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.sign_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.acos_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.top_k_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.asin_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.abs_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.ceil_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.shift_right_logical_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.exp_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.le_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.rem_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.asinh_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.ne_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.is_finite_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.lt_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.iota_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.select_n_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.atan_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.square_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.shift_left_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.acosh_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.clz_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.imag_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.clamp_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.transpose_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.broadcast_in_dim_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.integer_pow_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.reduce_or_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.le_to_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.nextafter_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.shift_right_arithmetic_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.argmin_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.argmax_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.div_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.gt_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.tan_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.not_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.reduce_and_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.pow_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.log1p_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.sort_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.dot_general_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.add_p", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.atan2_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.logistic_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.rev_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.copy_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.reduce_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.convert_element_type_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.squeeze_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.max_p", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.sin_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.and_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.reduce_precision_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.exp2_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.rsqrt_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.reduce_sum_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.eq_to_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.floor_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.lax.lax.gt", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.transpose", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.bitwise_or", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax._src.lax.lax.tile_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.lax.lax.rev", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.empty", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.lax.integer_pow", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.lax.lax.RaggedDotDimensionNumbers", "methods": [ { "kind": "function", "name": "RaggedDotDimensionNumbers.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "RaggedDotDimensionNumbers.dot_dimension_numbers", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RaggedDotDimensionNumbers.lhs_ragged_dimensions", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RaggedDotDimensionNumbers.rhs_group_dimensions", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 4, "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 3, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax._src.lax.lax.rsqrt", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.iota", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.rng_uniform", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.lax.argmax", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.lax.atan", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.convert_element_type", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.broadcast_in_dim", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.lax.lax.log1p", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.select_n", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.asinh", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.batch_matmul", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.lax.pad", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.lax.clamp", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.lax.shape_as_value", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.argmin", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.lax.sub", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.concatenate", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.sqrt", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.split", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.lax.sign", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.reciprocal", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.reshape", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.lax.lax.dce_sink", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.square", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.full_like", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "attr", "name": "jax._src.lax.lax.optimization_barrier_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "class", "name": "jax._src.lax.lax.RoundingMethod", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "RoundingMethod.AWAY_FROM_ZERO", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RoundingMethod.TO_NEAREST_EVEN", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax._src.lax.lax.rem", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.asin", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.dot_general", "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "class", "name": "jax._src.lax.lax.RandomAlgorithm", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "RandomAlgorithm.RNG_DEFAULT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RandomAlgorithm.RNG_THREE_FRY", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "RandomAlgorithm.RNG_PHILOX", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax._src.lax.lax.create_token", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.complex", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.abs", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.population_count", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.cosh", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.reduce_precision", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.lax.round", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.lax.lax.Tolerance", "methods": [ { "kind": "function", "name": "Tolerance.__init__", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Tolerance.atol", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Tolerance.rtol", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Tolerance.ulps", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 7, "n_typed": 3, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 3, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax._src.lax.lax.broadcast_shapes", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 2, "n_typable": 2, "n_params": 1, "n_function_overloads": 2, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.reduce_and", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.sort", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 2, "n_typable": 5, "n_params": 4, "n_function_overloads": 2, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.lax.lax.bitwise_and", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.exp2", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.squeeze", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.bitcast_convert_type", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.exp", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax._src.lax.lax.dtype", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax._src.lax.lax.bitwise_xor", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.ragged_dot_general", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.lax.lax.reduce_min", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.nextafter", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.is_finite", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.acos", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.add", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.cos", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.atan2", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.max", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.shift_left", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.broadcast_to_rank", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax._src.lax.lax.dce_sink_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "class", "name": "jax._src.lax.lax.DotAlgorithm", "methods": [ { "kind": "function", "name": "DotAlgorithm._convert_to_hlo_attr", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "DotAlgorithm.lhs_precision_type", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithm.rhs_precision_type", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithm.accumulation_type", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithm.lhs_component_count", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithm.rhs_component_count", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithm.num_primitive_operations", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithm.allow_imprecise_accumulation", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 3, "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 2, "n_attrs": 7, "n_properties": 0 }, { "kind": "function", "name": "jax._src.lax.lax.lt", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.pow", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.ne", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.acosh", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.reduce_or", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.rng_bit_generator", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.lax.lax.reduce_xor", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.le", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.reduce", "n_typed": 3, "n_any": 3, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.lax.lax.tanh", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.real", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.conj", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.full", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.lax.lax.imag", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.collapse", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.lax.reduce_prod", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.expand_dims", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax._src.lax.lax.split_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.lax.lax.composite", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.lax.cbrt", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.eq", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.padtype_to_pads", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.lax.lax.bitwise_not", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.tile", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.broadcast", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.lax.log", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.ceil", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax._src.lax.lax.DotDimensionNumbers", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax._src.lax.lax.dot", "n_typed": 6, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "class", "name": "jax._src.lax.lax.DotAlgorithmPreset", "methods": [ { "kind": "function", "name": "DotAlgorithmPreset.__repr__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DotAlgorithmPreset.__str__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "DotAlgorithmPreset.supported_output_types", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "DotAlgorithmPreset._convert_to_hlo_attr", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [ { "kind": "property", "name": "DotAlgorithmPreset.supported_lhs_types", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "DotAlgorithmPreset.supported_rhs_types", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "DotAlgorithmPreset.accumulation_type", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "DotAlgorithmPreset.DEFAULT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithmPreset.ANY_F8_ANY_F8_F32", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithmPreset.ANY_F8_ANY_F8_F32_FAST_ACCUM", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithmPreset.ANY_F8_ANY_F8_ANY", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithmPreset.ANY_F8_ANY_F8_ANY_FAST_ACCUM", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithmPreset.F16_F16_F16", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithmPreset.F16_F16_F32", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithmPreset.BF16_BF16_BF16", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithmPreset.BF16_BF16_F32", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithmPreset.BF16_BF16_F32_X3", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithmPreset.BF16_BF16_F32_X6", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithmPreset.BF16_BF16_F32_X9", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithmPreset.TF32_TF32_F32", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithmPreset.TF32_TF32_F32_X3", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithmPreset.F32_F32_F32", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "DotAlgorithmPreset.F64_F64_F64", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 11, "n_typed": 11, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 4, "n_method_overloads": 4, "n_method_params": 4, "n_attrs": 16, "n_properties": 3 }, { "kind": "function", "name": "jax._src.lax.lax.stop_gradient", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax._src.lax.lax.AccuracyMode", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "AccuracyMode.HIGHEST", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "AccuracyMode.DEFAULT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax._src.lax.lax.ragged_dot", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.lax.lax.optimization_barrier", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.min", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.neg", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.tan", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.reduce_sum", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.lax.reduce_max", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.shift_right_arithmetic", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.sin", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.select", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.lax.expm1", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.logistic", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.floor", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.after_all", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.mul", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.sort_key_val", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.lax.lax.top_k", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.lax.shift_right_logical", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.sinh", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.broadcasted_iota", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "class", "name": "jax._src.lax.lax.Precision", "methods": [ { "kind": "function", "name": "Precision._missing_", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Precision.__repr__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Precision.__str__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Precision.DEFAULT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Precision.HIGH", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Precision.HIGHEST", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 4, "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 1, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax._src.lax.lax.atanh", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.div", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.lax.clz", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.lax.ge", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax._src.lax.lax.PrecisionLike", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "type_ignores": [ { "kind": "type", "rules": [ "bad-assignment" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": [ "method-assign" ] } ], "name": "jax._src.lax.lax", "names": [ "jax._src.lax.lax.AccuracyMode", "jax._src.lax.lax.DotAlgorithm", "jax._src.lax.lax.DotAlgorithmPreset", "jax._src.lax.lax.DotDimensionNumbers", "jax._src.lax.lax.Precision", "jax._src.lax.lax.PrecisionLike", "jax._src.lax.lax.RaggedDotDimensionNumbers", "jax._src.lax.lax.RandomAlgorithm", "jax._src.lax.lax.RoundingMethod", "jax._src.lax.lax.Tolerance", "jax._src.lax.lax.abs", "jax._src.lax.lax.abs_p", "jax._src.lax.lax.acos", "jax._src.lax.lax.acos_p", "jax._src.lax.lax.acosh", "jax._src.lax.lax.acosh_p", "jax._src.lax.lax.add", "jax._src.lax.lax.add_p", "jax._src.lax.lax.after_all", "jax._src.lax.lax.after_all_p", "jax._src.lax.lax.and_p", "jax._src.lax.lax.argmax", "jax._src.lax.lax.argmax_p", "jax._src.lax.lax.argmin", "jax._src.lax.lax.argmin_p", "jax._src.lax.lax.asin", "jax._src.lax.lax.asin_p", "jax._src.lax.lax.asinh", "jax._src.lax.lax.asinh_p", "jax._src.lax.lax.atan", "jax._src.lax.lax.atan2", "jax._src.lax.lax.atan2_p", "jax._src.lax.lax.atan_p", "jax._src.lax.lax.atanh", "jax._src.lax.lax.atanh_p", "jax._src.lax.lax.batch_matmul", "jax._src.lax.lax.bitcast_convert_type", "jax._src.lax.lax.bitcast_convert_type_p", "jax._src.lax.lax.bitwise_and", "jax._src.lax.lax.bitwise_not", "jax._src.lax.lax.bitwise_or", "jax._src.lax.lax.bitwise_xor", "jax._src.lax.lax.broadcast", "jax._src.lax.lax.broadcast_in_dim", "jax._src.lax.lax.broadcast_in_dim_p", "jax._src.lax.lax.broadcast_shapes", "jax._src.lax.lax.broadcast_to_rank", "jax._src.lax.lax.broadcasted_iota", "jax._src.lax.lax.cbrt", "jax._src.lax.lax.cbrt_p", "jax._src.lax.lax.ceil", "jax._src.lax.lax.ceil_p", "jax._src.lax.lax.clamp", "jax._src.lax.lax.clamp_p", "jax._src.lax.lax.clz", "jax._src.lax.lax.clz_p", "jax._src.lax.lax.collapse", "jax._src.lax.lax.complex", "jax._src.lax.lax.complex_p", "jax._src.lax.lax.composite", "jax._src.lax.lax.concatenate", "jax._src.lax.lax.concatenate_p", "jax._src.lax.lax.conj", "jax._src.lax.lax.conj_p", "jax._src.lax.lax.convert_element_type", "jax._src.lax.lax.convert_element_type_p", "jax._src.lax.lax.copy_p", "jax._src.lax.lax.cos", "jax._src.lax.lax.cos_p", "jax._src.lax.lax.cosh", "jax._src.lax.lax.cosh_p", "jax._src.lax.lax.create_token", "jax._src.lax.lax.create_token_p", "jax._src.lax.lax.dce_sink", "jax._src.lax.lax.dce_sink_p", "jax._src.lax.lax.div", "jax._src.lax.lax.div_p", "jax._src.lax.lax.dot", "jax._src.lax.lax.dot_general", "jax._src.lax.lax.dot_general_p", "jax._src.lax.lax.dtype", "jax._src.lax.lax.empty", "jax._src.lax.lax.empty2_p", "jax._src.lax.lax.eq", "jax._src.lax.lax.eq_p", "jax._src.lax.lax.eq_to_p", "jax._src.lax.lax.exp", "jax._src.lax.lax.exp2", "jax._src.lax.lax.exp2_p", "jax._src.lax.lax.exp_p", "jax._src.lax.lax.expand_dims", "jax._src.lax.lax.expm1", "jax._src.lax.lax.expm1_p", "jax._src.lax.lax.floor", "jax._src.lax.lax.floor_p", "jax._src.lax.lax.full", "jax._src.lax.lax.full_like", "jax._src.lax.lax.ge", "jax._src.lax.lax.ge_p", "jax._src.lax.lax.gt", "jax._src.lax.lax.gt_p", "jax._src.lax.lax.imag", "jax._src.lax.lax.imag_p", "jax._src.lax.lax.integer_pow", "jax._src.lax.lax.integer_pow_p", "jax._src.lax.lax.iota", "jax._src.lax.lax.iota_p", "jax._src.lax.lax.is_finite", "jax._src.lax.lax.is_finite_p", "jax._src.lax.lax.le", "jax._src.lax.lax.le_p", "jax._src.lax.lax.le_to_p", "jax._src.lax.lax.log", "jax._src.lax.lax.log1p", "jax._src.lax.lax.log1p_p", "jax._src.lax.lax.log_p", "jax._src.lax.lax.logistic", "jax._src.lax.lax.logistic_p", "jax._src.lax.lax.lt", "jax._src.lax.lax.lt_p", "jax._src.lax.lax.lt_to_p", "jax._src.lax.lax.max", "jax._src.lax.lax.max_p", "jax._src.lax.lax.min", "jax._src.lax.lax.min_p", "jax._src.lax.lax.mul", "jax._src.lax.lax.mul_p", "jax._src.lax.lax.ne", "jax._src.lax.lax.ne_p", "jax._src.lax.lax.neg", "jax._src.lax.lax.neg_p", "jax._src.lax.lax.nextafter", "jax._src.lax.lax.nextafter_p", "jax._src.lax.lax.not_p", "jax._src.lax.lax.optimization_barrier", "jax._src.lax.lax.optimization_barrier_p", "jax._src.lax.lax.or_p", "jax._src.lax.lax.pad", "jax._src.lax.lax.pad_p", "jax._src.lax.lax.padtype_to_pads", "jax._src.lax.lax.population_count", "jax._src.lax.lax.population_count_p", "jax._src.lax.lax.pow", "jax._src.lax.lax.pow_p", "jax._src.lax.lax.ragged_dot", "jax._src.lax.lax.ragged_dot_general", "jax._src.lax.lax.real", "jax._src.lax.lax.real_p", "jax._src.lax.lax.reciprocal", "jax._src.lax.lax.reduce", "jax._src.lax.lax.reduce_and", "jax._src.lax.lax.reduce_and_p", "jax._src.lax.lax.reduce_max", "jax._src.lax.lax.reduce_max_p", "jax._src.lax.lax.reduce_min", "jax._src.lax.lax.reduce_min_p", "jax._src.lax.lax.reduce_or", "jax._src.lax.lax.reduce_or_p", "jax._src.lax.lax.reduce_p", "jax._src.lax.lax.reduce_precision", "jax._src.lax.lax.reduce_precision_p", "jax._src.lax.lax.reduce_prod", "jax._src.lax.lax.reduce_prod_p", "jax._src.lax.lax.reduce_sum", "jax._src.lax.lax.reduce_sum_p", "jax._src.lax.lax.reduce_xor", "jax._src.lax.lax.reduce_xor_p", "jax._src.lax.lax.rem", "jax._src.lax.lax.rem_p", "jax._src.lax.lax.reshape", "jax._src.lax.lax.reshape_p", "jax._src.lax.lax.rev", "jax._src.lax.lax.rev_p", "jax._src.lax.lax.rng_bit_generator", "jax._src.lax.lax.rng_bit_generator_p", "jax._src.lax.lax.rng_uniform", "jax._src.lax.lax.rng_uniform_p", "jax._src.lax.lax.round", "jax._src.lax.lax.round_p", "jax._src.lax.lax.rsqrt", "jax._src.lax.lax.rsqrt_p", "jax._src.lax.lax.select", "jax._src.lax.lax.select_n", "jax._src.lax.lax.select_n_p", "jax._src.lax.lax.shape_as_value", "jax._src.lax.lax.shift_left", "jax._src.lax.lax.shift_left_p", "jax._src.lax.lax.shift_right_arithmetic", "jax._src.lax.lax.shift_right_arithmetic_p", "jax._src.lax.lax.shift_right_logical", "jax._src.lax.lax.shift_right_logical_p", "jax._src.lax.lax.sign", "jax._src.lax.lax.sign_p", "jax._src.lax.lax.sin", "jax._src.lax.lax.sin_p", "jax._src.lax.lax.sinh", "jax._src.lax.lax.sinh_p", "jax._src.lax.lax.sort", "jax._src.lax.lax.sort_key_val", "jax._src.lax.lax.sort_p", "jax._src.lax.lax.split", "jax._src.lax.lax.split_p", "jax._src.lax.lax.sqrt", "jax._src.lax.lax.sqrt_p", "jax._src.lax.lax.square", "jax._src.lax.lax.square_p", "jax._src.lax.lax.squeeze", "jax._src.lax.lax.squeeze_p", "jax._src.lax.lax.stop_gradient", "jax._src.lax.lax.sub", "jax._src.lax.lax.sub_p", "jax._src.lax.lax.tan", "jax._src.lax.lax.tan_p", "jax._src.lax.lax.tanh", "jax._src.lax.lax.tanh_p", "jax._src.lax.lax.tile", "jax._src.lax.lax.tile_p", "jax._src.lax.lax.top_k", "jax._src.lax.lax.top_k_p", "jax._src.lax.lax.transpose", "jax._src.lax.lax.transpose_p", "jax._src.lax.lax.xor_p" ], "n_typable": 495, "n_typed": 345, "n_any": 3, "n_untyped": 147, "n_functions": 113, "n_function_overloads": 115, "n_function_params": 253, "n_methods": 10, "n_method_overloads": 10, "n_method_params": 13, "n_classes": 8, "n_attrs": 140, "n_properties": 3, "n_type_ignores": 3 }, { "path": "jax/_src/lax/linalg.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.lax.linalg.schur_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.linalg.cholesky_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.linalg.householder_product_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.linalg.eigh_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.linalg.svd_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.linalg.hessenberg_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.linalg.eig_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.linalg.tridiagonal_solve_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.linalg.lu_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.linalg.tridiagonal_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.linalg.triangular_solve_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.linalg.qr_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.lax.linalg.symmetric_product", "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "attr", "name": "jax._src.lax.linalg.cholesky_update_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.lax.linalg.lu_pivots_to_permutation", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.linalg.cholesky_update", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.linalg.householder_product", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax._src.lax.linalg.lu_pivots_to_permutation_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.lax.linalg.lu", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.linalg.eig", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.lax.linalg.eigh", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "class", "name": "jax._src.lax.linalg.EighImplementation", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "EighImplementation.QR", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "EighImplementation.JACOBI", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "EighImplementation.QDWH", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "class", "name": "jax._src.lax.linalg.EigImplementation", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "EigImplementation.CUSOLVER", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "EigImplementation.MAGMA", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "EigImplementation.LAPACK", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax._src.lax.linalg.qr", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 3, "n_typable": 5, "n_params": 4, "n_function_overloads": 3, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.lax.linalg.tridiagonal", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.lax.linalg.SvdAlgorithm", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "SvdAlgorithm.DEFAULT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SvdAlgorithm.QR", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SvdAlgorithm.JACOBI", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "SvdAlgorithm.POLAR", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 4, "n_properties": 0 }, { "kind": "function", "name": "jax._src.lax.linalg.tridiagonal_solve", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.lax.linalg.schur", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.lax.linalg.cholesky", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.linalg.svd", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 3, "n_typable": 6, "n_params": 5, "n_function_overloads": 3, "n_function_params": 5 }, { "kind": "attr", "name": "jax._src.lax.linalg.symmetric_product_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.lax.linalg.hessenberg", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.linalg.triangular_solve", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 } ], "type_ignores": [], "name": "jax._src.lax.linalg", "names": [ "jax._src.lax.linalg.EigImplementation", "jax._src.lax.linalg.EighImplementation", "jax._src.lax.linalg.SvdAlgorithm", "jax._src.lax.linalg.cholesky", "jax._src.lax.linalg.cholesky_p", "jax._src.lax.linalg.cholesky_update", "jax._src.lax.linalg.cholesky_update_p", "jax._src.lax.linalg.eig", "jax._src.lax.linalg.eig_p", "jax._src.lax.linalg.eigh", "jax._src.lax.linalg.eigh_p", "jax._src.lax.linalg.hessenberg", "jax._src.lax.linalg.hessenberg_p", "jax._src.lax.linalg.householder_product", "jax._src.lax.linalg.householder_product_p", "jax._src.lax.linalg.lu", "jax._src.lax.linalg.lu_p", "jax._src.lax.linalg.lu_pivots_to_permutation", "jax._src.lax.linalg.lu_pivots_to_permutation_p", "jax._src.lax.linalg.qr", "jax._src.lax.linalg.qr_p", "jax._src.lax.linalg.schur", "jax._src.lax.linalg.schur_p", "jax._src.lax.linalg.svd", "jax._src.lax.linalg.svd_p", "jax._src.lax.linalg.symmetric_product", "jax._src.lax.linalg.symmetric_product_p", "jax._src.lax.linalg.triangular_solve", "jax._src.lax.linalg.triangular_solve_p", "jax._src.lax.linalg.tridiagonal", "jax._src.lax.linalg.tridiagonal_p", "jax._src.lax.linalg.tridiagonal_solve", "jax._src.lax.linalg.tridiagonal_solve_p" ], "n_typable": 82, "n_typed": 66, "n_any": 0, "n_untyped": 16, "n_functions": 15, "n_function_overloads": 19, "n_function_params": 52, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 3, "n_attrs": 25, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/lax/ann.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.lax.ann.approx_top_k_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.lax.ann.approx_max_k", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.lax.ann.approx_min_k", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 } ], "type_ignores": [], "name": "jax._src.lax.ann", "names": [ "jax._src.lax.ann.approx_max_k", "jax._src.lax.ann.approx_min_k", "jax._src.lax.ann.approx_top_k_p" ], "n_typable": 15, "n_typed": 14, "n_any": 0, "n_untyped": 1, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 12, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/prng.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.prng.random_fold_in_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.prng.random_seed_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.prng.random_split_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.prng.threefry2x32_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.prng.random_bits_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.prng.threefry_2x32", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.prng.seed_with_impl", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.prng.random_seed", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax._src.prng.unsafe_rbg_prng_impl", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.prng.rbg_prng_impl", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.prng.threefry_prng_impl", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [ { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "override" ] } ], "name": "jax._src.prng", "names": [ "jax._src.prng.random_bits_p", "jax._src.prng.random_fold_in_p", "jax._src.prng.random_seed", "jax._src.prng.random_seed_p", "jax._src.prng.random_split_p", "jax._src.prng.rbg_prng_impl", "jax._src.prng.seed_with_impl", "jax._src.prng.threefry2x32_p", "jax._src.prng.threefry_2x32", "jax._src.prng.threefry_prng_impl", "jax._src.prng.unsafe_rbg_prng_impl" ], "n_typable": 17, "n_typed": 6, "n_any": 0, "n_untyped": 11, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 6, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 8, "n_properties": 0, "n_type_ignores": 9 }, { "path": "jax/_src/lax/special.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.lax.special.bessel_i1e_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.special.lgamma_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.special.erf_inv_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.special.igamma_grad_a_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.special.regularized_incomplete_beta_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.special.erf_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.special.igammac_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.special.bessel_i0e_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.special.erfc_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.special.zeta_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.special.igamma_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.special.polygamma_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.special.digamma_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.lax.special.zeta", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.special.erf", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.special.igamma_grad_a", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.special.bessel_i0e", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.special.polygamma", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.special.lgamma", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.special.random_gamma_grad", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.special.betainc", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.special.bessel_i1e", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.special.erfc", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.special.igammac", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.special.digamma", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.special.igamma", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.lax.special.erf_inv", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax._src.lax.special", "names": [ "jax._src.lax.special.bessel_i0e", "jax._src.lax.special.bessel_i0e_p", "jax._src.lax.special.bessel_i1e", "jax._src.lax.special.bessel_i1e_p", "jax._src.lax.special.betainc", "jax._src.lax.special.digamma", "jax._src.lax.special.digamma_p", "jax._src.lax.special.erf", "jax._src.lax.special.erf_inv", "jax._src.lax.special.erf_inv_p", "jax._src.lax.special.erf_p", "jax._src.lax.special.erfc", "jax._src.lax.special.erfc_p", "jax._src.lax.special.igamma", "jax._src.lax.special.igamma_grad_a", "jax._src.lax.special.igamma_grad_a_p", "jax._src.lax.special.igamma_p", "jax._src.lax.special.igammac", "jax._src.lax.special.igammac_p", "jax._src.lax.special.lgamma", "jax._src.lax.special.lgamma_p", "jax._src.lax.special.polygamma", "jax._src.lax.special.polygamma_p", "jax._src.lax.special.random_gamma_grad", "jax._src.lax.special.regularized_incomplete_beta_p", "jax._src.lax.special.zeta", "jax._src.lax.special.zeta_p" ], "n_typable": 50, "n_typed": 36, "n_any": 0, "n_untyped": 14, "n_functions": 14, "n_function_overloads": 14, "n_function_params": 23, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 13, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/lax/slicing.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.lax.slicing.scatter_mul_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.slicing.scatter_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.slicing.gather_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.slicing.scatter_add_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.slicing.dynamic_slice_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.slicing.scatter_max_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.slicing.dynamic_update_slice_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.slicing.scatter_min_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.slicing.slice_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.lax.slicing.index_in_dim", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.lax.slicing.index_take", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.slicing.scatter_sub", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.lax.slicing.scatter_apply", "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax._src.lax.slicing.dynamic_update_slice", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.lax.slicing.slice", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.lax.slicing.dynamic_index_in_dim", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.lax.slicing.dynamic_update_slice_in_dim", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "class", "name": "jax._src.lax.slicing.GatherDimensionNumbers", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "GatherDimensionNumbers.offset_dims", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GatherDimensionNumbers.collapsed_slice_dims", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GatherDimensionNumbers.start_index_map", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GatherDimensionNumbers.operand_batching_dims", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GatherDimensionNumbers.start_indices_batching_dims", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 5, "n_properties": 0 }, { "kind": "function", "name": "jax._src.lax.slicing.slice_in_dim", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.lax.slicing.scatter_max", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.lax.slicing.scatter", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "class", "name": "jax._src.lax.slicing.GatherScatterMode", "methods": [ { "kind": "function", "name": "GatherScatterMode.from_any", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "GatherScatterMode.CLIP", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GatherScatterMode.FILL_OR_DROP", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GatherScatterMode.PROMISE_IN_BOUNDS", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "GatherScatterMode.ONE_HOT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 2, "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 1, "n_attrs": 4, "n_properties": 0 }, { "kind": "function", "name": "jax._src.lax.slicing.scatter_min", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.lax.slicing.gather", "n_typed": 8, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax._src.lax.slicing.dynamic_slice", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.lax.slicing.dynamic_slice_in_dim", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "attr", "name": "jax._src.lax.slicing.scatter_sub_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.lax.slicing.scatter_add", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.lax.slicing.scatter_mul", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.lax.slicing.dynamic_update_index_in_dim", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "class", "name": "jax._src.lax.slicing.ScatterDimensionNumbers", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "ScatterDimensionNumbers.update_window_dims", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ScatterDimensionNumbers.inserted_window_dims", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ScatterDimensionNumbers.scatter_dims_to_operand_dims", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ScatterDimensionNumbers.operand_batching_dims", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ScatterDimensionNumbers.scatter_indices_batching_dims", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 5, "n_properties": 0 } ], "type_ignores": [], "name": "jax._src.lax.slicing", "names": [ "jax._src.lax.slicing.GatherDimensionNumbers", "jax._src.lax.slicing.GatherScatterMode", "jax._src.lax.slicing.ScatterDimensionNumbers", "jax._src.lax.slicing.dynamic_index_in_dim", "jax._src.lax.slicing.dynamic_slice", "jax._src.lax.slicing.dynamic_slice_in_dim", "jax._src.lax.slicing.dynamic_slice_p", "jax._src.lax.slicing.dynamic_update_index_in_dim", "jax._src.lax.slicing.dynamic_update_slice", "jax._src.lax.slicing.dynamic_update_slice_in_dim", "jax._src.lax.slicing.dynamic_update_slice_p", "jax._src.lax.slicing.gather", "jax._src.lax.slicing.gather_p", "jax._src.lax.slicing.index_in_dim", "jax._src.lax.slicing.index_take", "jax._src.lax.slicing.scatter", "jax._src.lax.slicing.scatter_add", "jax._src.lax.slicing.scatter_add_p", "jax._src.lax.slicing.scatter_apply", "jax._src.lax.slicing.scatter_max", "jax._src.lax.slicing.scatter_max_p", "jax._src.lax.slicing.scatter_min", "jax._src.lax.slicing.scatter_min_p", "jax._src.lax.slicing.scatter_mul", "jax._src.lax.slicing.scatter_mul_p", "jax._src.lax.slicing.scatter_p", "jax._src.lax.slicing.scatter_sub", "jax._src.lax.slicing.scatter_sub_p", "jax._src.lax.slicing.slice", "jax._src.lax.slicing.slice_in_dim", "jax._src.lax.slicing.slice_p" ], "n_typable": 132, "n_typed": 121, "n_any": 0, "n_untyped": 11, "n_functions": 18, "n_function_overloads": 18, "n_function_params": 102, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 1, "n_classes": 3, "n_attrs": 24, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/lax/control_flow/loops.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.lax.control_flow.loops.scan_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.control_flow.loops.cumlogsumexp_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.control_flow.loops.cumprod_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.control_flow.loops.cumsum_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.control_flow.loops.cummax_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.control_flow.loops.cummin_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.control_flow.loops.while_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.lax.control_flow.loops.scan", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.lax.control_flow.loops.fori_loop", "n_typed": 1, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.lax.control_flow.loops.associative_scan", "n_typed": 3, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.lax.control_flow.loops.map", "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.control_flow.loops.cumsum", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.control_flow.loops.cummax", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.control_flow.loops.cumprod", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.control_flow.loops.cumlogsumexp", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.control_flow.loops.while_loop", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.control_flow.loops.cummin", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [ { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null } ], "name": "jax._src.lax.control_flow.loops", "names": [ "jax._src.lax.control_flow.loops.associative_scan", "jax._src.lax.control_flow.loops.cumlogsumexp", "jax._src.lax.control_flow.loops.cumlogsumexp_p", "jax._src.lax.control_flow.loops.cummax", "jax._src.lax.control_flow.loops.cummax_p", "jax._src.lax.control_flow.loops.cummin", "jax._src.lax.control_flow.loops.cummin_p", "jax._src.lax.control_flow.loops.cumprod", "jax._src.lax.control_flow.loops.cumprod_p", "jax._src.lax.control_flow.loops.cumsum", "jax._src.lax.control_flow.loops.cumsum_p", "jax._src.lax.control_flow.loops.fori_loop", "jax._src.lax.control_flow.loops.map", "jax._src.lax.control_flow.loops.scan", "jax._src.lax.control_flow.loops.scan_p", "jax._src.lax.control_flow.loops.while_loop", "jax._src.lax.control_flow.loops.while_p" ], "n_typable": 54, "n_typed": 37, "n_any": 0, "n_untyped": 17, "n_functions": 10, "n_function_overloads": 10, "n_function_params": 37, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 7, "n_properties": 0, "n_type_ignores": 4 }, { "path": "jax/_src/lax/parallel.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.lax.parallel.pmin_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.parallel.ragged_all_to_all_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.parallel.all_gather_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.parallel.pmax_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.parallel.psum_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.parallel.ppermute_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.parallel.all_to_all_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.parallel.axis_index_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.lax.parallel.pshuffle", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.parallel.psum", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.parallel.ppermute", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.parallel.psend", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.parallel.ragged_all_to_all", "n_typed": 0, "n_any": 0, "n_untyped": 9, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax._src.lax.parallel.pcast", "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.parallel.precv", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.lax.parallel.pmean", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.parallel.axis_index", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.parallel.axis_size", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.parallel.pmin", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.parallel.pswapaxes", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.lax.parallel.all_gather_start", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.lax.parallel.all_to_all", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.lax.parallel.all_gather_done", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.parallel.all_gather", "n_typed": 1, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.lax.parallel.pmax", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.parallel.psum_scatter", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.lax.parallel.pbroadcast", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [ { "kind": "type", "rules": null }, { "kind": "type", "rules": null } ], "name": "jax._src.lax.parallel", "names": [ "jax._src.lax.parallel.all_gather", "jax._src.lax.parallel.all_gather_done", "jax._src.lax.parallel.all_gather_p", "jax._src.lax.parallel.all_gather_start", "jax._src.lax.parallel.all_to_all", "jax._src.lax.parallel.all_to_all_p", "jax._src.lax.parallel.axis_index", "jax._src.lax.parallel.axis_index_p", "jax._src.lax.parallel.axis_size", "jax._src.lax.parallel.pbroadcast", "jax._src.lax.parallel.pcast", "jax._src.lax.parallel.pmax", "jax._src.lax.parallel.pmax_p", "jax._src.lax.parallel.pmean", "jax._src.lax.parallel.pmin", "jax._src.lax.parallel.pmin_p", "jax._src.lax.parallel.ppermute", "jax._src.lax.parallel.ppermute_p", "jax._src.lax.parallel.precv", "jax._src.lax.parallel.psend", "jax._src.lax.parallel.pshuffle", "jax._src.lax.parallel.psum", "jax._src.lax.parallel.psum_p", "jax._src.lax.parallel.psum_scatter", "jax._src.lax.parallel.pswapaxes", "jax._src.lax.parallel.ragged_all_to_all", "jax._src.lax.parallel.ragged_all_to_all_p" ], "n_typable": 94, "n_typed": 6, "n_any": 0, "n_untyped": 88, "n_functions": 19, "n_function_overloads": 19, "n_function_params": 67, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 8, "n_properties": 0, "n_type_ignores": 2 }, { "path": "jax/_src/lax/fft.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.lax.fft.fft_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "class", "name": "jax._src.lax.fft.FftType", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "FftType.FFT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "FftType.IFFT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "FftType.RFFT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "FftType.IRFFT", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 4, "n_properties": 0 }, { "kind": "function", "name": "jax._src.lax.fft.fft", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.lax.fft", "names": [ "jax._src.lax.fft.FftType", "jax._src.lax.fft.fft", "jax._src.lax.fft.fft_p" ], "n_typable": 5, "n_typed": 2, "n_any": 0, "n_untyped": 3, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 3, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 5, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/lax/control_flow/solves.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.lax.control_flow.solves.linear_solve_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.lax.control_flow.solves.custom_linear_solve", "n_typed": 3, "n_any": 1, "n_untyped": 3, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.lax.control_flow.solves.custom_root", "n_typed": 3, "n_any": 1, "n_untyped": 2, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "type_ignores": [], "name": "jax._src.lax.control_flow.solves", "names": [ "jax._src.lax.control_flow.solves.custom_linear_solve", "jax._src.lax.control_flow.solves.custom_root", "jax._src.lax.control_flow.solves.linear_solve_p" ], "n_typable": 14, "n_typed": 6, "n_any": 2, "n_untyped": 6, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 11, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/random.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.random.random_gamma_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.random.PRNGKey", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.random.gumbel", "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.random.randint", "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.random.laplace", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.random.clone", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.random.f", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.random.cauchy", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.random.multinomial", "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.random.multivariate_normal", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.random.key", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.random.key_data", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.random.geometric", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.random.beta", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.random.wald", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.random.rademacher", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.random.rayleigh", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.random.lognormal", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.random.poisson", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.random.double_sided_maxwell", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.random.key_impl", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.random.wrap_key_data", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.random.chisquare", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.random.maxwell", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.random.triangular", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.random.bits", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.random.pareto", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.random.orthogonal", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.random.categorical", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.random.bernoulli", "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.random.weibull_min", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.random.t", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.random.uniform", "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.random.gamma", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.random.choice", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.random.generalized_normal", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.random.exponential", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.random.truncated_normal", "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.random.ball", "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.random.split", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.random.loggamma", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.random.dirichlet", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.random.binomial", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.random.logistic", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.random.fold_in", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.random.permutation", "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.random.normal", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax._src.random", "names": [ "jax._src.random.PRNGKey", "jax._src.random.ball", "jax._src.random.bernoulli", "jax._src.random.beta", "jax._src.random.binomial", "jax._src.random.bits", "jax._src.random.categorical", "jax._src.random.cauchy", "jax._src.random.chisquare", "jax._src.random.choice", "jax._src.random.clone", "jax._src.random.dirichlet", "jax._src.random.double_sided_maxwell", "jax._src.random.exponential", "jax._src.random.f", "jax._src.random.fold_in", "jax._src.random.gamma", "jax._src.random.generalized_normal", "jax._src.random.geometric", "jax._src.random.gumbel", "jax._src.random.key", "jax._src.random.key_data", "jax._src.random.key_impl", "jax._src.random.laplace", "jax._src.random.loggamma", "jax._src.random.logistic", "jax._src.random.lognormal", "jax._src.random.maxwell", "jax._src.random.multinomial", "jax._src.random.multivariate_normal", "jax._src.random.normal", "jax._src.random.orthogonal", "jax._src.random.pareto", "jax._src.random.permutation", "jax._src.random.poisson", "jax._src.random.rademacher", "jax._src.random.randint", "jax._src.random.random_gamma_p", "jax._src.random.rayleigh", "jax._src.random.split", "jax._src.random.t", "jax._src.random.triangular", "jax._src.random.truncated_normal", "jax._src.random.uniform", "jax._src.random.wald", "jax._src.random.weibull_min", "jax._src.random.wrap_key_data" ], "n_typable": 233, "n_typed": 219, "n_any": 0, "n_untyped": 14, "n_functions": 46, "n_function_overloads": 46, "n_function_params": 186, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/dispatch.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.dispatch.device_put_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.dispatch.apply_primitive", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.dispatch", "names": [ "jax._src.dispatch.apply_primitive", "jax._src.dispatch.device_put_p" ], "n_typable": 5, "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 3, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/lax/windowed_reductions.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.lax.windowed_reductions.reduce_window_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.windowed_reductions.reduce_window_min_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.windowed_reductions.select_and_gather_add_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.windowed_reductions.select_and_scatter_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.windowed_reductions.reduce_window_max_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.windowed_reductions.reduce_window_sum_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.lax.windowed_reductions.select_and_scatter_add_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.lax.windowed_reductions.reduce_window", "n_typed": 6, "n_any": 3, "n_untyped": 0, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax._src.lax.windowed_reductions.reduce_window_shape_tuple", "n_typed": 0, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 } ], "type_ignores": [], "name": "jax._src.lax.windowed_reductions", "names": [ "jax._src.lax.windowed_reductions.reduce_window", "jax._src.lax.windowed_reductions.reduce_window_max_p", "jax._src.lax.windowed_reductions.reduce_window_min_p", "jax._src.lax.windowed_reductions.reduce_window_p", "jax._src.lax.windowed_reductions.reduce_window_shape_tuple", "jax._src.lax.windowed_reductions.reduce_window_sum_p", "jax._src.lax.windowed_reductions.select_and_gather_add_p", "jax._src.lax.windowed_reductions.select_and_scatter_add_p", "jax._src.lax.windowed_reductions.select_and_scatter_p" ], "n_typable": 23, "n_typed": 6, "n_any": 3, "n_untyped": 14, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 14, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 7, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/lax/convolution.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.lax.convolution.conv_general_dilated_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.lax.convolution.conv_with_general_padding", "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax._src.lax.convolution.conv_general_shape_tuple", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.lax.convolution.conv_transpose_shape_tuple", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.lax.convolution.conv", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.lax.convolution.conv_general_dilated", "n_typed": 13, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 13, "n_params": 12, "n_function_overloads": 1, "n_function_params": 12 }, { "kind": "function", "name": "jax._src.lax.convolution.conv_dimension_numbers", "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.convolution.conv_general_permutations", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.lax.convolution.conv_transpose", "n_typed": 11, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 11, "n_params": 10, "n_function_overloads": 1, "n_function_params": 10 }, { "kind": "function", "name": "jax._src.lax.convolution.conv_shape_tuple", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "class", "name": "jax._src.lax.convolution.ConvDimensionNumbers", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "ConvDimensionNumbers.lhs_spec", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ConvDimensionNumbers.rhs_spec", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ConvDimensionNumbers.out_spec", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "attr", "name": "jax._src.lax.convolution.ConvGeneralDilatedDimensionNumbers", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "type_ignores": [], "name": "jax._src.lax.convolution", "names": [ "jax._src.lax.convolution.ConvDimensionNumbers", "jax._src.lax.convolution.ConvGeneralDilatedDimensionNumbers", "jax._src.lax.convolution.conv", "jax._src.lax.convolution.conv_dimension_numbers", "jax._src.lax.convolution.conv_general_dilated", "jax._src.lax.convolution.conv_general_dilated_p", "jax._src.lax.convolution.conv_general_permutations", "jax._src.lax.convolution.conv_general_shape_tuple", "jax._src.lax.convolution.conv_shape_tuple", "jax._src.lax.convolution.conv_transpose", "jax._src.lax.convolution.conv_transpose_shape_tuple", "jax._src.lax.convolution.conv_with_general_padding" ], "n_typable": 66, "n_typed": 42, "n_any": 0, "n_untyped": 24, "n_functions": 9, "n_function_overloads": 9, "n_function_params": 55, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 5, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/lax/control_flow/conditionals.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.lax.control_flow.conditionals.cond_p", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.lax.control_flow.conditionals.cond", "n_typed": 2, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.lax.control_flow.conditionals.platform_dependent", "n_typed": 2, "n_any": 1, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.lax.control_flow.conditionals.switch", "n_typed": 1, "n_any": 2, "n_untyped": 2, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [ { "kind": "type", "rules": null } ], "name": "jax._src.lax.control_flow.conditionals", "names": [ "jax._src.lax.control_flow.conditionals.cond", "jax._src.lax.control_flow.conditionals.cond_p", "jax._src.lax.control_flow.conditionals.platform_dependent", "jax._src.lax.control_flow.conditionals.switch" ], "n_typable": 16, "n_typed": 5, "n_any": 3, "n_untyped": 8, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 12, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 1 }, { "path": "jax/extend/linear_util.py", "symbol_reports": [ { "kind": "function", "name": "jax.extend.linear_util.wrap_init", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax.extend.linear_util", "names": [ "jax.extend.linear_util.wrap_init" ], "n_typable": 4, "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 3, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/interpreters/mlir.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.interpreters.mlir.lower_with_sharding_in_types", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.interpreters.mlir.flatten_ir_values", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.interpreters.mlir.merge_mlir_modules", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "class", "name": "jax._src.interpreters.mlir.LoweringResult", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "LoweringResult.module", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringResult.keepalive", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringResult.host_callbacks", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringResult.shape_poly_state", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 4, "n_properties": 0 }, { "kind": "function", "name": "jax._src.interpreters.mlir.lower_jaxpr_to_fun", "n_typed": 21, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 21, "n_params": 20, "n_function_overloads": 1, "n_function_params": 20 }, { "kind": "function", "name": "jax._src.interpreters.mlir.dtype_to_ir_type", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax._src.interpreters.mlir.AxisContext", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax._src.interpreters.mlir.unflatten_ir_values_like_types", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax._src.interpreters.mlir.RECV_FROM_HOST_TYPE", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax._src.interpreters.mlir.register_constant_handler", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.interpreters.mlir.aval_to_ir_types", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.interpreters.mlir.lower_jaxpr_to_module", "n_typed": 23, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 23, "n_params": 22, "n_function_overloads": 1, "n_function_params": 22 }, { "kind": "attr", "name": "jax._src.interpreters.mlir.SEND_TO_HOST_TYPE", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax._src.interpreters.mlir.ir_attribute", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax._src.interpreters.mlir.ir_type_handlers", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "class", "name": "jax._src.interpreters.mlir.TokenSet", "methods": [ { "kind": "function", "name": "TokenSet.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "TokenSet.__len__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "TokenSet.get", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TokenSet.create", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TokenSet.items", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "TokenSet.effects", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "TokenSet.subset", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "TokenSet.update_tokens", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "TokenSet._tokens", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 15, "n_typed": 11, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 8, "n_method_overloads": 8, "n_method_params": 6, "n_attrs": 1, "n_properties": 0 }, { "kind": "function", "name": "jax._src.interpreters.mlir.shape_tensor", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax._src.interpreters.mlir.LoweringRule", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax._src.interpreters.mlir.flatten_ir_types", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.interpreters.mlir.module_to_bytecode", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.interpreters.mlir.module_to_string", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.interpreters.mlir.ShapePolyLoweringState", "methods": [ { "kind": "function", "name": "ShapePolyLoweringState.__init__", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "ShapePolyLoweringState.dim_vars", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "ShapePolyLoweringState.uses_dim_vars", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "ShapePolyLoweringState.has_platform_index_argument", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 6, "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 2, "n_attrs": 3, "n_properties": 0 }, { "kind": "function", "name": "jax._src.interpreters.mlir.i32_attr", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax._src.interpreters.mlir.ModuleContext", "methods": [ { "kind": "function", "name": "ModuleContext.__init__", "n_typed": 15, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 17, "n_params": 16, "n_function_overloads": 1, "n_function_params": 16 }, { "kind": "function", "name": "ModuleContext.get_backend", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ModuleContext.new_channel", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "ModuleContext.add_host_callback", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ModuleContext.add_keepalive", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ModuleContext.replace", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "ModuleContext.axis_env", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "ModuleContext.context", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ModuleContext.module", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ModuleContext.ip", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ModuleContext.symbol_table", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ModuleContext.platforms", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ModuleContext.backend", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ModuleContext.axis_context", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ModuleContext.keepalives", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ModuleContext.channel_iterator", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ModuleContext.host_callbacks", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ModuleContext.shape_poly_state", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ModuleContext.all_default_mem_kind", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ModuleContext.lowering_cache", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ModuleContext.cached_primitive_lowerings", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ModuleContext.traceback_caches", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ModuleContext.lowering_parameters", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 27, "n_typed": 21, "n_any": 2, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 6, "n_method_overloads": 6, "n_method_params": 20, "n_attrs": 16, "n_properties": 1 }, { "kind": "function", "name": "jax._src.interpreters.mlir.lower_fun", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.interpreters.mlir.make_ir_context", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "class", "name": "jax._src.interpreters.mlir.LoweringParameters", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "LoweringParameters.override_lowering_rules", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringParameters.global_constant_computation", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringParameters.for_export", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringParameters.export_ignore_forward_compatibility", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringParameters.hoist_constants_as_args", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 5, "n_properties": 0 }, { "kind": "function", "name": "jax._src.interpreters.mlir.dense_int_elements", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax._src.interpreters.mlir.ConstantHandler", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.interpreters.mlir.LoweringRuleContext", "methods": [ { "kind": "function", "name": "LoweringRuleContext.set_tokens_out", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "LoweringRuleContext.replace", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "LoweringRuleContext.is_forward_compat", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "LoweringRuleContext.module_context", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringRuleContext.name_stack", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringRuleContext.traceback", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringRuleContext.primitive", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringRuleContext.avals_in", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringRuleContext.avals_out", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringRuleContext.tokens_in", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringRuleContext.tokens_out", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringRuleContext.const_lowering", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringRuleContext.axis_size_env", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringRuleContext.dim_var_values", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringRuleContext.jaxpr_eqn_ctx", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "LoweringRuleContext.platforms", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 5, "n_typed": 2, "n_any": 0, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 2, "n_attrs": 13, "n_properties": 0 }, { "kind": "function", "name": "jax._src.interpreters.mlir.ir_constant", "n_typed": 3, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.interpreters.mlir.core_call_lowering", "n_typed": 2, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "attr", "name": "jax._src.interpreters.mlir.DEVICE_TO_DEVICE_TYPE", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax._src.interpreters.mlir.i64_attr", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.interpreters.mlir.register_lowering", "n_typed": 5, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.interpreters.mlir.aval_to_ir_type", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.interpreters.mlir.jaxpr_subcomp", "n_typed": 10, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 } ], "type_ignores": [ { "kind": "type", "rules": null }, { "kind": "type", "rules": null } ], "name": "jax._src.interpreters.mlir", "names": [ "jax._src.interpreters.mlir.AxisContext", "jax._src.interpreters.mlir.ConstantHandler", "jax._src.interpreters.mlir.DEVICE_TO_DEVICE_TYPE", "jax._src.interpreters.mlir.LoweringParameters", "jax._src.interpreters.mlir.LoweringResult", "jax._src.interpreters.mlir.LoweringRule", "jax._src.interpreters.mlir.LoweringRuleContext", "jax._src.interpreters.mlir.ModuleContext", "jax._src.interpreters.mlir.RECV_FROM_HOST_TYPE", "jax._src.interpreters.mlir.SEND_TO_HOST_TYPE", "jax._src.interpreters.mlir.ShapePolyLoweringState", "jax._src.interpreters.mlir.TokenSet", "jax._src.interpreters.mlir.aval_to_ir_type", "jax._src.interpreters.mlir.aval_to_ir_types", "jax._src.interpreters.mlir.core_call_lowering", "jax._src.interpreters.mlir.dense_int_elements", "jax._src.interpreters.mlir.dtype_to_ir_type", "jax._src.interpreters.mlir.flatten_ir_types", "jax._src.interpreters.mlir.flatten_ir_values", "jax._src.interpreters.mlir.i32_attr", "jax._src.interpreters.mlir.i64_attr", "jax._src.interpreters.mlir.ir_attribute", "jax._src.interpreters.mlir.ir_constant", "jax._src.interpreters.mlir.ir_type_handlers", "jax._src.interpreters.mlir.jaxpr_subcomp", "jax._src.interpreters.mlir.lower_fun", "jax._src.interpreters.mlir.lower_jaxpr_to_fun", "jax._src.interpreters.mlir.lower_jaxpr_to_module", "jax._src.interpreters.mlir.lower_with_sharding_in_types", "jax._src.interpreters.mlir.make_ir_context", "jax._src.interpreters.mlir.merge_mlir_modules", "jax._src.interpreters.mlir.module_to_bytecode", "jax._src.interpreters.mlir.module_to_string", "jax._src.interpreters.mlir.register_constant_handler", "jax._src.interpreters.mlir.register_lowering", "jax._src.interpreters.mlir.shape_tensor", "jax._src.interpreters.mlir.unflatten_ir_values_like_types" ], "n_typable": 170, "n_typed": 137, "n_any": 5, "n_untyped": 28, "n_functions": 24, "n_function_overloads": 24, "n_function_params": 91, "n_methods": 18, "n_method_overloads": 18, "n_method_params": 30, "n_classes": 8, "n_attrs": 47, "n_properties": 1, "n_type_ignores": 2 }, { "path": "jax/_src/extend/random.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.extend.random.define_prng_impl", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 } ], "type_ignores": [], "name": "jax._src.extend.random", "names": [ "jax._src.extend.random.define_prng_impl" ], "n_typable": 8, "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 7, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/extend/sharding.py", "symbol_reports": [ { "kind": "function", "name": "jax.extend.sharding.get_serialized_proto_from_hlo_sharding", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.extend.sharding.get_hlo_sharding_from_serialized_proto", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.extend.sharding.get_op_sharding_from_serialized_proto", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.extend.sharding", "names": [ "jax.extend.sharding.get_hlo_sharding_from_serialized_proto", "jax.extend.sharding.get_op_sharding_from_serialized_proto", "jax.extend.sharding.get_serialized_proto_from_hlo_sharding" ], "n_typable": 6, "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 3, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/source_info_util.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.source_info_util.summarize", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.source_info_util.new_source_info", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.source_info_util.new_name_stack", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax._src.source_info_util.NameStack", "methods": [ { "kind": "function", "name": "NameStack.extend", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "NameStack.transform", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "NameStack.__getitem__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "NameStack.__len__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "NameStack.__add__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "NameStack.__radd__", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "NameStack.__str__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "NameStack.stack", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 12, "n_typed": 11, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 5, "n_attrs": 1, "n_properties": 0 }, { "kind": "function", "name": "jax._src.source_info_util.register_exclusion", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.source_info_util.current", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "class", "name": "jax._src.source_info_util.SourceInfo", "methods": [ { "kind": "function", "name": "SourceInfo.__init__", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "SourceInfo.replace", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "SourceInfo.traceback", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "SourceInfo.name_stack", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 8, "n_typed": 7, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 4, "n_attrs": 2, "n_properties": 0 }, { "kind": "function", "name": "jax._src.source_info_util.current_name_stack", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.source_info_util.reset_name_stack", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "type_ignores": [], "name": "jax._src.source_info_util", "names": [ "jax._src.source_info_util.NameStack", "jax._src.source_info_util.SourceInfo", "jax._src.source_info_util.current", "jax._src.source_info_util.current_name_stack", "jax._src.source_info_util.new_name_stack", "jax._src.source_info_util.new_source_info", "jax._src.source_info_util.register_exclusion", "jax._src.source_info_util.reset_name_stack", "jax._src.source_info_util.summarize" ], "n_typable": 31, "n_typed": 27, "n_any": 0, "n_untyped": 4, "n_functions": 7, "n_function_overloads": 7, "n_function_params": 4, "n_methods": 9, "n_method_overloads": 9, "n_method_params": 9, "n_classes": 2, "n_attrs": 3, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/extend/source_info_util.py", "symbol_reports": [ { "kind": "attr", "name": "jax.extend.source_info_util.transform_name_stack", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.extend.source_info_util.extend_name_stack", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.extend.source_info_util.user_context", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.extend.source_info_util.set_name_stack", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [], "name": "jax.extend.source_info_util", "names": [ "jax.extend.source_info_util.extend_name_stack", "jax.extend.source_info_util.set_name_stack", "jax.extend.source_info_util.transform_name_stack", "jax.extend.source_info_util.user_context" ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 4, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/ffi.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.ffi.register_ffi_target", "n_typed": 4, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.ffi.register_ffi_type", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.ffi.ffi_call", "n_typed": 11, "n_any": 0, "n_untyped": 0, "n_overloads": 2, "n_typable": 11, "n_params": 10, "n_function_overloads": 2, "n_function_params": 10 }, { "kind": "function", "name": "jax._src.ffi.register_ffi_target_as_batch_partitionable", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.ffi.register_ffi_type_id", "n_typed": 3, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.ffi.ffi_lowering", "n_typed": 6, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.ffi.include_dir", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.ffi.pycapsule", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.ffi.build_ffi_lowering_function", "n_typed": 6, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 } ], "type_ignores": [], "name": "jax._src.ffi", "names": [ "jax._src.ffi.build_ffi_lowering_function", "jax._src.ffi.ffi_call", "jax._src.ffi.ffi_lowering", "jax._src.ffi.include_dir", "jax._src.ffi.pycapsule", "jax._src.ffi.register_ffi_target", "jax._src.ffi.register_ffi_target_as_batch_partitionable", "jax._src.ffi.register_ffi_type", "jax._src.ffi.register_ffi_type_id" ], "n_typable": 44, "n_typed": 37, "n_any": 5, "n_untyped": 2, "n_functions": 9, "n_function_overloads": 10, "n_function_params": 35, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/flatten_util.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.flatten_util.ravel_pytree", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax._src.flatten_util", "names": [ "jax._src.flatten_util.ravel_pytree" ], "n_typable": 2, "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 1, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/image/scale.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.image.scale.resize", "n_typed": 3, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.image.scale.scale_and_translate", "n_typed": 4, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "class", "name": "jax._src.image.scale.ResizeMethod", "methods": [ { "kind": "function", "name": "ResizeMethod.from_string", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "ResizeMethod.NEAREST", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ResizeMethod.LINEAR", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ResizeMethod.LANCZOS3", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ResizeMethod.LANCZOS5", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ResizeMethod.CUBIC", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 2, "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 1, "n_attrs": 5, "n_properties": 0 } ], "type_ignores": [], "name": "jax._src.image.scale", "names": [ "jax._src.image.scale.ResizeMethod", "jax._src.image.scale.resize", "jax._src.image.scale.scale_and_translate" ], "n_typable": 17, "n_typed": 8, "n_any": 0, "n_untyped": 9, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 13, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 1, "n_classes": 1, "n_attrs": 5, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/interpreters/batching.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.interpreters.batching.defbroadcasting", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.interpreters.batching.register_vmappable", "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "attr", "name": "jax._src.interpreters.batching.axis_primitive_batchers", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.interpreters.batching.bdim_at_front", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "attr", "name": "jax._src.interpreters.batching.primitive_batchers", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "function", "name": "jax._src.interpreters.batching.broadcast", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "attr", "name": "jax._src.interpreters.batching.fancy_primitive_batchers", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax._src.interpreters.batching.unregister_vmappable", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax._src.interpreters.batching.not_mapped", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax._src.interpreters.batching.defreducer", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.interpreters.batching.defvectorized", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [ { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-param-name-override" ] }, { "kind": "pyrefly", "rules": [ "bad-param-name-override" ] } ], "name": "jax._src.interpreters.batching", "names": [ "jax._src.interpreters.batching.axis_primitive_batchers", "jax._src.interpreters.batching.bdim_at_front", "jax._src.interpreters.batching.broadcast", "jax._src.interpreters.batching.defbroadcasting", "jax._src.interpreters.batching.defreducer", "jax._src.interpreters.batching.defvectorized", "jax._src.interpreters.batching.fancy_primitive_batchers", "jax._src.interpreters.batching.not_mapped", "jax._src.interpreters.batching.primitive_batchers", "jax._src.interpreters.batching.register_vmappable", "jax._src.interpreters.batching.unregister_vmappable" ], "n_typable": 28, "n_typed": 9, "n_any": 0, "n_untyped": 19, "n_functions": 7, "n_function_overloads": 7, "n_function_params": 18, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 4, "n_properties": 0, "n_type_ignores": 3 }, { "path": "jax/interpreters/mlir.py", "symbol_reports": [ { "kind": "attr", "name": "jax.interpreters.mlir.MeshAxisName", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.interpreters.mlir.Token", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.interpreters.mlir.token_type", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.interpreters.mlir.Value", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.interpreters.mlir.ir", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.interpreters.mlir.dense_int_array", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [], "name": "jax.interpreters.mlir", "names": [ "jax.interpreters.mlir.MeshAxisName", "jax.interpreters.mlir.Token", "jax.interpreters.mlir.Value", "jax.interpreters.mlir.dense_int_array", "jax.interpreters.mlir.ir", "jax.interpreters.mlir.token_type" ], "n_typable": 3, "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 6, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/mesh.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.mesh.Mesh", "methods": [ { "kind": "function", "name": "Mesh.__new__", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Mesh.__reduce__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Mesh.__eq__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Mesh.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Mesh.__setattr__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Mesh.__enter__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Mesh.__exit__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Mesh.update", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Mesh._local_mesh", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Mesh.__str__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Mesh.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "Mesh.shape", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Mesh.shape_tuple", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Mesh.axis_sizes", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Mesh.size", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Mesh.empty", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Mesh.is_multi_process", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Mesh.local_mesh", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Mesh.device_ids", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Mesh._local_devices_set", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Mesh._flat_devices_tuple", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Mesh._internal_device_list", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Mesh._flat_devices_set", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Mesh._repr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Mesh.local_devices", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Mesh.abstract_mesh", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "Mesh.devices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Mesh.axis_names", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 41, "n_typed": 6, "n_any": 0, "n_untyped": 35, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 11, "n_method_overloads": 11, "n_method_params": 13, "n_attrs": 2, "n_properties": 15 }, { "kind": "class", "name": "jax._src.mesh.use_abstract_mesh", "methods": [ { "kind": "function", "name": "use_abstract_mesh.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "use_abstract_mesh.__enter__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "use_abstract_mesh.__exit__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "use_abstract_mesh.mesh", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 8, "n_typed": 1, "n_any": 0, "n_untyped": 7, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 4, "n_attrs": 1, "n_properties": 0 }, { "kind": "class", "name": "jax._src.mesh.AbstractMesh", "methods": [ { "kind": "function", "name": "AbstractMesh.__init__", "n_typed": 3, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "AbstractMesh.__hash__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AbstractMesh.__eq__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AbstractMesh.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AbstractMesh.update", "n_typed": 0, "n_any": 0, "n_untyped": 5, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "AbstractMesh.update_axis_types", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "AbstractMesh.__enter__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AbstractMesh.__exit__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [ { "kind": "property", "name": "AbstractMesh.shape", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "AbstractMesh.shape_tuple", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "AbstractMesh._internal_device_list", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "AbstractMesh.empty", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "AbstractMesh.abstract_mesh", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "AbstractMesh.devices", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "AbstractMesh.device_ids", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "AbstractMesh.is_multi_process", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "AbstractMesh.local_devices", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "AbstractMesh.local_mesh", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "AbstractMesh.axis_sizes", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "AbstractMesh.abstract_device", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "AbstractMesh.size", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 34, "n_typed": 4, "n_any": 0, "n_untyped": 30, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 8, "n_method_overloads": 8, "n_method_params": 13, "n_attrs": 3, "n_properties": 10 }, { "kind": "function", "name": "jax._src.mesh.get_abstract_mesh", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "class", "name": "jax._src.mesh.AxisType", "methods": [ { "kind": "function", "name": "AxisType.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "AxisType.Auto", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "AxisType.Explicit", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "AxisType.Manual", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 }, { "kind": "class", "name": "jax._src.mesh.AbstractDevice", "methods": [ { "kind": "function", "name": "AbstractDevice.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "AbstractDevice._repr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "AbstractDevice.device_kind", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "AbstractDevice.num_cores", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 2, "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 0, "n_attrs": 2, "n_properties": 0 } ], "type_ignores": [ { "kind": "type", "rules": null } ], "name": "jax._src.mesh", "names": [ "jax._src.mesh.AbstractDevice", "jax._src.mesh.AbstractMesh", "jax._src.mesh.AxisType", "jax._src.mesh.Mesh", "jax._src.mesh.get_abstract_mesh", "jax._src.mesh.use_abstract_mesh" ], "n_typable": 87, "n_typed": 12, "n_any": 0, "n_untyped": 75, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 0, "n_methods": 25, "n_method_overloads": 25, "n_method_params": 30, "n_classes": 5, "n_attrs": 11, "n_properties": 25, "n_type_ignores": 1 }, { "path": "jax/interpreters/partial_eval.py", "symbol_reports": [ { "kind": "attr", "name": "jax.interpreters.partial_eval.Val", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [], "name": "jax.interpreters.partial_eval", "names": [ "jax.interpreters.partial_eval.Val" ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/interpreters/xla.py", "symbol_reports": [ { "kind": "attr", "name": "jax.interpreters.xla.Backend", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [], "name": "jax.interpreters.xla", "names": [ "jax.interpreters.xla.Backend" ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/lax/other.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.lax.other.conv_general_dilated_patches", "n_typed": 10, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "jax._src.lax.other.conv_general_dilated_local", "n_typed": 10, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 } ], "type_ignores": [], "name": "jax._src.lax.other", "names": [ "jax._src.lax.other.conv_general_dilated_local", "jax._src.lax.other.conv_general_dilated_patches" ], "n_typable": 20, "n_typed": 20, "n_any": 0, "n_untyped": 0, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 18, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/lax/scaled_dot.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.lax.scaled_dot.scaled_dot", "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 } ], "type_ignores": [], "name": "jax._src.lax.scaled_dot", "names": [ "jax._src.lax.scaled_dot.scaled_dot" ], "n_typable": 7, "n_typed": 6, "n_any": 0, "n_untyped": 1, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 6, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/tpu/linalg/qdwh.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.tpu.linalg.qdwh.qdwh", "n_typed": 4, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "type_ignores": [], "name": "jax._src.tpu.linalg.qdwh", "names": [ "jax._src.tpu.linalg.qdwh.qdwh" ], "n_typable": 6, "n_typed": 4, "n_any": 0, "n_untyped": 2, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 5, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/memory.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.memory.Space", "methods": [ { "kind": "function", "name": "Space.__repr__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Space.Device", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Space.Host", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Space.Any", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_attrs": 3, "n_properties": 0 } ], "type_ignores": [], "name": "jax._src.memory", "names": [ "jax._src.memory.Space" ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 0, "n_classes": 1, "n_attrs": 3, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/monitoring.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.monitoring.record_event", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.monitoring.record_event_duration_secs", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.monitoring.unregister_event_duration_listener", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.monitoring.unregister_event_listener", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.monitoring.register_event_listener", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.monitoring.record_event_time_span", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.monitoring.clear_event_listeners", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.monitoring.register_event_duration_secs_listener", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.monitoring.unregister_scalar_listener", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.monitoring.register_event_time_span_listener", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.monitoring.record_scalar", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.monitoring.unregister_event_time_span_listener", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.monitoring.register_scalar_listener", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax._src.monitoring", "names": [ "jax._src.monitoring.clear_event_listeners", "jax._src.monitoring.record_event", "jax._src.monitoring.record_event_duration_secs", "jax._src.monitoring.record_event_time_span", "jax._src.monitoring.record_scalar", "jax._src.monitoring.register_event_duration_secs_listener", "jax._src.monitoring.register_event_listener", "jax._src.monitoring.register_event_time_span_listener", "jax._src.monitoring.register_scalar_listener", "jax._src.monitoring.unregister_event_duration_listener", "jax._src.monitoring.unregister_event_listener", "jax._src.monitoring.unregister_event_time_span_listener", "jax._src.monitoring.unregister_scalar_listener" ], "n_typable": 33, "n_typed": 32, "n_any": 0, "n_untyped": 1, "n_functions": 13, "n_function_overloads": 13, "n_function_params": 20, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/nn/__init__.pyi", "symbol_reports": [ { "kind": "function", "name": "jax.nn.elu", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.nn.sigmoid", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.nn.hard_silu", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.nn.get_scaled_dot_general_config", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.nn.logmeanexp", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.nn.softmax", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.nn.celu", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.nn.logsumexp", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 3, "n_typable": 7, "n_params": 6, "n_function_overloads": 3, "n_function_params": 6 }, { "kind": "function", "name": "jax.nn.glu", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.nn.log_softmax", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.nn.leaky_relu", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.nn.sparse_sigmoid", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.nn.log1mexp", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.nn.relu6", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.nn.one_hot", "n_typed": 3, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.nn.standardize", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.nn.silu", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.nn.dot_product_attention", "n_typed": 13, "n_any": 0, "n_untyped": 0, "n_overloads": 2, "n_typable": 13, "n_params": 12, "n_function_overloads": 2, "n_function_params": 12 }, { "kind": "function", "name": "jax.nn.relu", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.nn.hard_tanh", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.nn.hard_swish", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.nn.mish", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.nn.hard_sigmoid", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.nn.gelu", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.nn.softplus", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.nn.sparse_plus", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.nn.squareplus", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.nn.log_sigmoid", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.nn.tanh", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.nn.selu", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.nn.Axis", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.nn.scaled_matmul", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.nn.scaled_dot_general", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.nn.swish", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.nn.identity", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.nn.soft_sign", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax.nn", "names": [ "jax.nn.Axis", "jax.nn.celu", "jax.nn.dot_product_attention", "jax.nn.elu", "jax.nn.gelu", "jax.nn.get_scaled_dot_general_config", "jax.nn.glu", "jax.nn.hard_sigmoid", "jax.nn.hard_silu", "jax.nn.hard_swish", "jax.nn.hard_tanh", "jax.nn.identity", "jax.nn.leaky_relu", "jax.nn.log1mexp", "jax.nn.log_sigmoid", "jax.nn.log_softmax", "jax.nn.logmeanexp", "jax.nn.logsumexp", "jax.nn.mish", "jax.nn.one_hot", "jax.nn.relu", "jax.nn.relu6", "jax.nn.scaled_dot_general", "jax.nn.scaled_matmul", "jax.nn.selu", "jax.nn.sigmoid", "jax.nn.silu", "jax.nn.soft_sign", "jax.nn.softmax", "jax.nn.softplus", "jax.nn.sparse_plus", "jax.nn.sparse_sigmoid", "jax.nn.squareplus", "jax.nn.standardize", "jax.nn.swish", "jax.nn.tanh" ], "n_typable": 117, "n_typed": 115, "n_any": 2, "n_untyped": 0, "n_functions": 35, "n_function_overloads": 38, "n_function_params": 82, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/nn/initializers.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.nn.initializers.zeros", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.nn.initializers.lecun_uniform", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.nn.initializers.lecun_normal", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.nn.initializers.glorot_uniform", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.nn.initializers.orthogonal", "n_typed": 3, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax._src.nn.initializers.Initializer", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax._src.nn.initializers.constant", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.nn.initializers.uniform", "n_typed": 2, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.nn.initializers.he_uniform", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.nn.initializers.he_normal", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.nn.initializers.truncated_normal", "n_typed": 2, "n_any": 3, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.nn.initializers.glorot_normal", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.nn.initializers.ones", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.nn.initializers.variance_scaling", "n_typed": 7, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.nn.initializers.normal", "n_typed": 2, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.nn.initializers.delta_orthogonal", "n_typed": 3, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.nn.initializers", "names": [ "jax._src.nn.initializers.Initializer", "jax._src.nn.initializers.constant", "jax._src.nn.initializers.delta_orthogonal", "jax._src.nn.initializers.glorot_normal", "jax._src.nn.initializers.glorot_uniform", "jax._src.nn.initializers.he_normal", "jax._src.nn.initializers.he_uniform", "jax._src.nn.initializers.lecun_normal", "jax._src.nn.initializers.lecun_uniform", "jax._src.nn.initializers.normal", "jax._src.nn.initializers.ones", "jax._src.nn.initializers.orthogonal", "jax._src.nn.initializers.truncated_normal", "jax._src.nn.initializers.uniform", "jax._src.nn.initializers.variance_scaling", "jax._src.nn.initializers.zeros" ], "n_typable": 70, "n_typed": 62, "n_any": 8, "n_untyped": 0, "n_functions": 15, "n_function_overloads": 15, "n_function_params": 55, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/nn/initializers.py", "symbol_reports": [ { "kind": "attr", "name": "jax.nn.initializers.kaiming_uniform", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.nn.initializers.xavier_normal", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.nn.initializers.xavier_uniform", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax.nn.initializers.kaiming_normal", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [], "name": "jax.nn.initializers", "names": [ "jax.nn.initializers.kaiming_normal", "jax.nn.initializers.kaiming_uniform", "jax.nn.initializers.xavier_normal", "jax.nn.initializers.xavier_uniform" ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 4, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/numpy/__init__.pyi", "symbol_reports": [ { "kind": "function", "name": "jax.numpy.atleast_2d", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 3, "n_typable": 4, "n_params": 3, "n_function_overloads": 3, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.rollaxis", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.divide", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.resize", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.trunc", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.put_along_axis", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.numpy.sort_complex", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.std", "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.numpy.degrees", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.transpose", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.diff", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "attr", "name": "jax.numpy.bitwise_or", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.tril", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.numpy.uint2", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.median", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.kaiser", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.polydiv", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.shape", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.arctanh", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.iscomplexobj", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.percentile", "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.numpy.isnan", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.numpy.pi", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.eye", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.left_shift", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.bartlett", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.fix", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.append", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax.numpy.euler_gamma", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.poly", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.packbits", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.astype", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.unique_values", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax.numpy.complex128", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.result_type", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.diag_indices_from", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.numpy.PadValueLike", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.unravel_index", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.from_dlpack", "n_typed": 3, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.arccos", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.diagonal", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.not_equal", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.nanmax", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.numpy.dstack", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.numpy.bfloat16", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.numpy.cdouble", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.diag_indices", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.empty", "n_typed": 4, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.unwrap", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.ptp", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.count_nonzero", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.rint", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.__array_namespace_info__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax.numpy.interp", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.numpy.size", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.fromfile", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.log2", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.negative", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.bitwise_invert", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.iscomplex", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.identity", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.permute_dims", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.any", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.amax", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.numpy.argmax", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.atan", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.broadcast_arrays", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.issubdtype", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.nancumsum", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.average", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 3, "n_typable": 6, "n_params": 5, "n_function_overloads": 3, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.log1p", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.setdiff1d", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.ones_like", "n_typed": 5, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.take_along_axis", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.asinh", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.greater_equal", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.isscalar", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.nanmin", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.numpy.vdot", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.pad", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.trapezoid", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.matrix_transpose", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "class", "name": "jax.numpy.BinaryUfunc", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "attr", "name": "jax.numpy.int2", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.nansum", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "attr", "name": "jax.numpy.newaxis", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax.numpy.kron", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.fmin", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.polyadd", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.argmin", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "attr", "name": "jax.numpy.multiply", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.isdtype", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.fabs", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.numpy.float8_e5m2fnuz", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.einsum", "n_typed": 8, "n_any": 0, "n_untyped": 2, "n_overloads": 3, "n_typable": 10, "n_params": 9, "n_function_overloads": 3, "n_function_params": 9 }, { "kind": "function", "name": "jax.numpy.flip", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.roots", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.setxor1d", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.concatenate", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.sqrt", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.split", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.convolve", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.bitwise_count", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.numpy.logical_xor", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.clip", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.numpy.ix_", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.sign", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.numpy.maximum", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.reciprocal", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.where", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 3, "n_typable": 6, "n_params": 5, "n_function_overloads": 3, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.conjugate", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.reshape", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.all", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.vander", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.square", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.full_like", "n_typed": 5, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "attr", "name": "jax.numpy.float32", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.asin", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.digitize", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.nanargmax", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.lexsort", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.cumulative_sum", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.histogram_bin_edges", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.indices", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 3, "n_typable": 4, "n_params": 3, "n_function_overloads": 3, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.dsplit", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.numpy.float8_e4m3", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.nanquantile", "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "attr", "name": "jax.numpy.uint64", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.nanprod", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "attr", "name": "jax.numpy.single", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.numpy.float_", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.numpy.mgrid", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.cumsum", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.abs", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.divmod", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.right_shift", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.frexp", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.numpy.float8_e4m3fn", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.einsum_path", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 3, "n_typable": 5, "n_params": 4, "n_function_overloads": 3, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.piecewise", "n_typed": 4, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "attr", "name": "jax.numpy.uint8", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.nanmean", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.numpy.logspace", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax.numpy.correlate", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.moveaxis", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.cosh", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.mod", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.nanvar", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax.numpy.round", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.ldexp", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.vstack", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.numpy.float4_e2m1fn", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.mean", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.numpy.broadcast_shapes", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 2, "n_typable": 2, "n_params": 1, "n_function_overloads": 2, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.bincount", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.repeat", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.polysub", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.sort", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "attr", "name": "jax.numpy.bitwise_and", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.trim_zeros", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.ediff1d", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.allclose", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.hypot", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.exp2", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.vecdot", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.radians", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.numpy.int16", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.greater", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.isin", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.squeeze", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.triu", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.exp", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.amin", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "attr", "name": "jax.numpy.int8", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.numpy.bitwise_xor", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.array_equiv", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.power", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.fromfunction", "n_typed": 3, "n_any": 1, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "attr", "name": "jax.numpy.int32", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.invert", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.arctan2", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.delete", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.tril_indices_from", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.array_split", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.nextafter", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.union1d", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.acos", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.flipud", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.numpy.float16", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.numpy.uint4", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.cumulative_prod", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.argsort", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "attr", "name": "jax.numpy.add", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.copy", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.array", "n_typed": 7, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax.numpy.compress", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.numpy.cos", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.polyint", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.atan2", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.triu_indices", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.gcd", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.arange", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.numpy.polyval", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.max", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.numpy.absolute", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.nonzero", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.rad2deg", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.isneginf", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.numpy.logical_and", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.atleast_1d", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 3, "n_typable": 4, "n_params": 3, "n_function_overloads": 3, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.pow", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.numpy.float8_e4m3fnuz", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.linspace", "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_overloads": 4, "n_typable": 9, "n_params": 8, "n_function_overloads": 4, "n_function_params": 8 }, { "kind": "attr", "name": "jax.numpy.int_", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.numpy.ComplexWarning", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.flatnonzero", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.nanstd", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax.numpy.vectorize", "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.signbit", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.rot90", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.arcsin", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.acosh", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.numpy.csingle", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.float_power", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.numpy.logaddexp", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.meshgrid", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.searchsorted", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.bitwise_left_shift", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.lcm", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.column_stack", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.concat", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.numpy.minimum", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.logical_not", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.numpy.float8_e4m3b11fnuz", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.true_divide", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.copysign", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.isinf", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.argpartition", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.isrealobj", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.numpy.uint32", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.sinc", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.zeros_like", "n_typed": 5, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.tanh", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.intersect1d", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.numpy.arctan", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.cumprod", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.fromstring", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.cross", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.numpy.floor_divide", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.numpy.uint16", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.angle", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.real", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.inner", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.apply_along_axis", "n_typed": 4, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.conj", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.full", "n_typed": 4, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.imag", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.tri", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.prod", "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.numpy.cov", "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.numpy.triu_indices_from", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.numpy.complex_", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.fliplr", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.histogram2d", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "attr", "name": "jax.numpy.float8_e3m4", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.hsplit", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.equal", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.log10", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.expand_dims", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.tensordot", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.numpy.less", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.polymul", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax.numpy.nan", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.cbrt", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.numpy.complex64", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.apply_over_axes", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax.numpy.logaddexp2", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.bitwise_not", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.matmul", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.tile", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax.numpy.ufunc", "methods": [ { "kind": "function", "name": "ufunc.__init__", "n_typed": 10, "n_any": 1, "n_untyped": 1, "n_overloads": 1, "n_typable": 12, "n_params": 11, "n_function_overloads": 1, "n_function_params": 11 }, { "kind": "function", "name": "ufunc.__call__", "n_typed": 1, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "ufunc.reduce", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "ufunc.accumulate", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "ufunc.at", "n_typed": 4, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "ufunc.reduceat", "n_typed": 5, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "ufunc.outer", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [ { "kind": "property", "name": "ufunc.nin", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "ufunc.nout", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "ufunc.nargs", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "ufunc.identity", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [], "n_typable": 45, "n_typed": 40, "n_any": 4, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 34, "n_attrs": 0, "n_properties": 4 }, { "kind": "function", "name": "jax.numpy.heaviside", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.asarray", "n_typed": 6, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.numpy.histogram", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.log", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.vsplit", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.sum", "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.numpy.ravel_multi_index", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.put", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.ceil", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.mask_indices", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.tril_indices", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.dot", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.frompyfunc", "n_typed": 4, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.quantile", "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.numpy.nanpercentile", "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.numpy.ravel", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.bitwise_right_shift", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.extract", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.roll", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.positive", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.place", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.gradient", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.fmax", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.min", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.numpy.frombuffer", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.swapaxes", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "attr", "name": "jax.numpy.float8_e5m2", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.insert", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.fmod", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.hstack", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.arcsinh", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.empty_like", "n_typed": 4, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.stack", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.unstack", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.fill_diagonal", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.trace", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "attr", "name": "jax.numpy.r_", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.tan", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.vecmat", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.numpy.uint", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.unique", "n_typed": 9, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "jax.numpy.var", "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.numpy.block", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.arccosh", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.unique_all", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.sin", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.unpackbits", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "attr", "name": "jax.numpy.inf", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.nanmedian", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.select", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.expm1", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.numpy.ogrid", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.floor", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.isposinf", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.isfinite", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.geomspace", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax.numpy.take", "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax.numpy.blackman", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.matvec", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.numpy.c_", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.array_equal", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.unique_inverse", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.outer", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.load", "n_typed": 2, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.around", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.spacing", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.polyder", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.sinh", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.unique_counts", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.less_equal", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.modf", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.numpy.logical_or", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.atanh", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.isreal", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.diagflat", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.zeros", "n_typed": 4, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "attr", "name": "jax.numpy.float8_e8m0fnu", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.numpy.bool_", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.isclose", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.deg2rad", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.nan_to_num", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax.numpy.atleast_3d", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 3, "n_typable": 4, "n_params": 3, "n_function_overloads": 3, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.polyfit", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax.numpy.argwhere", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.ndim", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.numpy.e", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.numpy.bool", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax.numpy.int64", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.partition", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.choose", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.corrcoef", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "attr", "name": "jax.numpy.float64", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.hamming", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.numpy.subtract", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.nancumprod", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.hanning", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax.numpy.int4", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.diag", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "attr", "name": "jax.numpy.__array_api_version__", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.broadcast_to", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax.numpy.histogramdd", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "attr", "name": "jax.numpy.double", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax.numpy.remainder", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.fromiter", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax.numpy.nanargmin", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax.numpy.i0", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax.numpy.ones", "n_typed": 4, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax.numpy", "names": [ "jax.numpy.BinaryUfunc", "jax.numpy.ComplexWarning", "jax.numpy.PadValueLike", "jax.numpy.__array_api_version__", "jax.numpy.__array_namespace_info__", "jax.numpy.abs", "jax.numpy.absolute", "jax.numpy.acos", "jax.numpy.acosh", "jax.numpy.add", "jax.numpy.all", "jax.numpy.allclose", "jax.numpy.amax", "jax.numpy.amin", "jax.numpy.angle", "jax.numpy.any", "jax.numpy.append", "jax.numpy.apply_along_axis", "jax.numpy.apply_over_axes", "jax.numpy.arange", "jax.numpy.arccos", "jax.numpy.arccosh", "jax.numpy.arcsin", "jax.numpy.arcsinh", "jax.numpy.arctan", "jax.numpy.arctan2", "jax.numpy.arctanh", "jax.numpy.argmax", "jax.numpy.argmin", "jax.numpy.argpartition", "jax.numpy.argsort", "jax.numpy.argwhere", "jax.numpy.around", "jax.numpy.array", "jax.numpy.array_equal", "jax.numpy.array_equiv", "jax.numpy.array_split", "jax.numpy.asarray", "jax.numpy.asin", "jax.numpy.asinh", "jax.numpy.astype", "jax.numpy.atan", "jax.numpy.atan2", "jax.numpy.atanh", "jax.numpy.atleast_1d", "jax.numpy.atleast_2d", "jax.numpy.atleast_3d", "jax.numpy.average", "jax.numpy.bartlett", "jax.numpy.bfloat16", "jax.numpy.bincount", "jax.numpy.bitwise_and", "jax.numpy.bitwise_count", "jax.numpy.bitwise_invert", "jax.numpy.bitwise_left_shift", "jax.numpy.bitwise_not", "jax.numpy.bitwise_or", "jax.numpy.bitwise_right_shift", "jax.numpy.bitwise_xor", "jax.numpy.blackman", "jax.numpy.block", "jax.numpy.bool", "jax.numpy.bool_", "jax.numpy.broadcast_arrays", "jax.numpy.broadcast_shapes", "jax.numpy.broadcast_to", "jax.numpy.c_", "jax.numpy.cbrt", "jax.numpy.cdouble", "jax.numpy.ceil", "jax.numpy.choose", "jax.numpy.clip", "jax.numpy.column_stack", "jax.numpy.complex128", "jax.numpy.complex64", "jax.numpy.complex_", "jax.numpy.compress", "jax.numpy.concat", "jax.numpy.concatenate", "jax.numpy.conj", "jax.numpy.conjugate", "jax.numpy.convolve", "jax.numpy.copy", "jax.numpy.copysign", "jax.numpy.corrcoef", "jax.numpy.correlate", "jax.numpy.cos", "jax.numpy.cosh", "jax.numpy.count_nonzero", "jax.numpy.cov", "jax.numpy.cross", "jax.numpy.csingle", "jax.numpy.cumprod", "jax.numpy.cumsum", "jax.numpy.cumulative_prod", "jax.numpy.cumulative_sum", "jax.numpy.deg2rad", "jax.numpy.degrees", "jax.numpy.delete", "jax.numpy.diag", "jax.numpy.diag_indices", "jax.numpy.diag_indices_from", "jax.numpy.diagflat", "jax.numpy.diagonal", "jax.numpy.diff", "jax.numpy.digitize", "jax.numpy.divide", "jax.numpy.divmod", "jax.numpy.dot", "jax.numpy.double", "jax.numpy.dsplit", "jax.numpy.dstack", "jax.numpy.e", "jax.numpy.ediff1d", "jax.numpy.einsum", "jax.numpy.einsum_path", "jax.numpy.empty", "jax.numpy.empty_like", "jax.numpy.equal", "jax.numpy.euler_gamma", "jax.numpy.exp", "jax.numpy.exp2", "jax.numpy.expand_dims", "jax.numpy.expm1", "jax.numpy.extract", "jax.numpy.eye", "jax.numpy.fabs", "jax.numpy.fill_diagonal", "jax.numpy.fix", "jax.numpy.flatnonzero", "jax.numpy.flip", "jax.numpy.fliplr", "jax.numpy.flipud", "jax.numpy.float16", "jax.numpy.float32", "jax.numpy.float4_e2m1fn", "jax.numpy.float64", "jax.numpy.float8_e3m4", "jax.numpy.float8_e4m3", "jax.numpy.float8_e4m3b11fnuz", "jax.numpy.float8_e4m3fn", "jax.numpy.float8_e4m3fnuz", "jax.numpy.float8_e5m2", "jax.numpy.float8_e5m2fnuz", "jax.numpy.float8_e8m0fnu", "jax.numpy.float_", "jax.numpy.float_power", "jax.numpy.floor", "jax.numpy.floor_divide", "jax.numpy.fmax", "jax.numpy.fmin", "jax.numpy.fmod", "jax.numpy.frexp", "jax.numpy.from_dlpack", "jax.numpy.frombuffer", "jax.numpy.fromfile", "jax.numpy.fromfunction", "jax.numpy.fromiter", "jax.numpy.frompyfunc", "jax.numpy.fromstring", "jax.numpy.full", "jax.numpy.full_like", "jax.numpy.gcd", "jax.numpy.geomspace", "jax.numpy.gradient", "jax.numpy.greater", "jax.numpy.greater_equal", "jax.numpy.hamming", "jax.numpy.hanning", "jax.numpy.heaviside", "jax.numpy.histogram", "jax.numpy.histogram2d", "jax.numpy.histogram_bin_edges", "jax.numpy.histogramdd", "jax.numpy.hsplit", "jax.numpy.hstack", "jax.numpy.hypot", "jax.numpy.i0", "jax.numpy.identity", "jax.numpy.imag", "jax.numpy.indices", "jax.numpy.inf", "jax.numpy.inner", "jax.numpy.insert", "jax.numpy.int16", "jax.numpy.int2", "jax.numpy.int32", "jax.numpy.int4", "jax.numpy.int64", "jax.numpy.int8", "jax.numpy.int_", "jax.numpy.interp", "jax.numpy.intersect1d", "jax.numpy.invert", "jax.numpy.isclose", "jax.numpy.iscomplex", "jax.numpy.iscomplexobj", "jax.numpy.isdtype", "jax.numpy.isfinite", "jax.numpy.isin", "jax.numpy.isinf", "jax.numpy.isnan", "jax.numpy.isneginf", "jax.numpy.isposinf", "jax.numpy.isreal", "jax.numpy.isrealobj", "jax.numpy.isscalar", "jax.numpy.issubdtype", "jax.numpy.ix_", "jax.numpy.kaiser", "jax.numpy.kron", "jax.numpy.lcm", "jax.numpy.ldexp", "jax.numpy.left_shift", "jax.numpy.less", "jax.numpy.less_equal", "jax.numpy.lexsort", "jax.numpy.linspace", "jax.numpy.load", "jax.numpy.log", "jax.numpy.log10", "jax.numpy.log1p", "jax.numpy.log2", "jax.numpy.logaddexp", "jax.numpy.logaddexp2", "jax.numpy.logical_and", "jax.numpy.logical_not", "jax.numpy.logical_or", "jax.numpy.logical_xor", "jax.numpy.logspace", "jax.numpy.mask_indices", "jax.numpy.matmul", "jax.numpy.matrix_transpose", "jax.numpy.matvec", "jax.numpy.max", "jax.numpy.maximum", "jax.numpy.mean", "jax.numpy.median", "jax.numpy.meshgrid", "jax.numpy.mgrid", "jax.numpy.min", "jax.numpy.minimum", "jax.numpy.mod", "jax.numpy.modf", "jax.numpy.moveaxis", "jax.numpy.multiply", "jax.numpy.nan", "jax.numpy.nan_to_num", "jax.numpy.nanargmax", "jax.numpy.nanargmin", "jax.numpy.nancumprod", "jax.numpy.nancumsum", "jax.numpy.nanmax", "jax.numpy.nanmean", "jax.numpy.nanmedian", "jax.numpy.nanmin", "jax.numpy.nanpercentile", "jax.numpy.nanprod", "jax.numpy.nanquantile", "jax.numpy.nanstd", "jax.numpy.nansum", "jax.numpy.nanvar", "jax.numpy.ndim", "jax.numpy.negative", "jax.numpy.newaxis", "jax.numpy.nextafter", "jax.numpy.nonzero", "jax.numpy.not_equal", "jax.numpy.ogrid", "jax.numpy.ones", "jax.numpy.ones_like", "jax.numpy.outer", "jax.numpy.packbits", "jax.numpy.pad", "jax.numpy.partition", "jax.numpy.percentile", "jax.numpy.permute_dims", "jax.numpy.pi", "jax.numpy.piecewise", "jax.numpy.place", "jax.numpy.poly", "jax.numpy.polyadd", "jax.numpy.polyder", "jax.numpy.polydiv", "jax.numpy.polyfit", "jax.numpy.polyint", "jax.numpy.polymul", "jax.numpy.polysub", "jax.numpy.polyval", "jax.numpy.positive", "jax.numpy.pow", "jax.numpy.power", "jax.numpy.prod", "jax.numpy.ptp", "jax.numpy.put", "jax.numpy.put_along_axis", "jax.numpy.quantile", "jax.numpy.r_", "jax.numpy.rad2deg", "jax.numpy.radians", "jax.numpy.ravel", "jax.numpy.ravel_multi_index", "jax.numpy.real", "jax.numpy.reciprocal", "jax.numpy.remainder", "jax.numpy.repeat", "jax.numpy.reshape", "jax.numpy.resize", "jax.numpy.result_type", "jax.numpy.right_shift", "jax.numpy.rint", "jax.numpy.roll", "jax.numpy.rollaxis", "jax.numpy.roots", "jax.numpy.rot90", "jax.numpy.round", "jax.numpy.searchsorted", "jax.numpy.select", "jax.numpy.setdiff1d", "jax.numpy.setxor1d", "jax.numpy.shape", "jax.numpy.sign", "jax.numpy.signbit", "jax.numpy.sin", "jax.numpy.sinc", "jax.numpy.single", "jax.numpy.sinh", "jax.numpy.size", "jax.numpy.sort", "jax.numpy.sort_complex", "jax.numpy.spacing", "jax.numpy.split", "jax.numpy.sqrt", "jax.numpy.square", "jax.numpy.squeeze", "jax.numpy.stack", "jax.numpy.std", "jax.numpy.subtract", "jax.numpy.sum", "jax.numpy.swapaxes", "jax.numpy.take", "jax.numpy.take_along_axis", "jax.numpy.tan", "jax.numpy.tanh", "jax.numpy.tensordot", "jax.numpy.tile", "jax.numpy.trace", "jax.numpy.transpose", "jax.numpy.trapezoid", "jax.numpy.tri", "jax.numpy.tril", "jax.numpy.tril_indices", "jax.numpy.tril_indices_from", "jax.numpy.trim_zeros", "jax.numpy.triu", "jax.numpy.triu_indices", "jax.numpy.triu_indices_from", "jax.numpy.true_divide", "jax.numpy.trunc", "jax.numpy.ufunc", "jax.numpy.uint", "jax.numpy.uint16", "jax.numpy.uint2", "jax.numpy.uint32", "jax.numpy.uint4", "jax.numpy.uint64", "jax.numpy.uint8", "jax.numpy.union1d", "jax.numpy.unique", "jax.numpy.unique_all", "jax.numpy.unique_counts", "jax.numpy.unique_inverse", "jax.numpy.unique_values", "jax.numpy.unpackbits", "jax.numpy.unravel_index", "jax.numpy.unstack", "jax.numpy.unwrap", "jax.numpy.vander", "jax.numpy.var", "jax.numpy.vdot", "jax.numpy.vecdot", "jax.numpy.vecmat", "jax.numpy.vectorize", "jax.numpy.vsplit", "jax.numpy.vstack", "jax.numpy.where", "jax.numpy.zeros", "jax.numpy.zeros_like" ], "n_typable": 1421, "n_typed": 1339, "n_any": 60, "n_untyped": 22, "n_functions": 323, "n_function_overloads": 343, "n_function_params": 991, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 34, "n_classes": 2, "n_attrs": 63, "n_properties": 4, "n_type_ignores": 0 }, { "path": "jax/_src/numpy/fft.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.numpy.fft.irfft2", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.fft.irfftn", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.fft.fft", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.fft.rfft", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.fft.irfft", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.fft.ifftn", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.fft.rfft2", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.fft.fftfreq", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.fft.ifft2", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.fft.rfftn", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.fft.fftn", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.fft.fft2", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.fft.hfft", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.fft.ifftshift", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.numpy.fft.fftshift", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.numpy.fft.ihfft", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.fft.rfftfreq", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.fft.ifft", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax._src.numpy.fft", "names": [ "jax._src.numpy.fft.fft", "jax._src.numpy.fft.fft2", "jax._src.numpy.fft.fftfreq", "jax._src.numpy.fft.fftn", "jax._src.numpy.fft.fftshift", "jax._src.numpy.fft.hfft", "jax._src.numpy.fft.ifft", "jax._src.numpy.fft.ifft2", "jax._src.numpy.fft.ifftn", "jax._src.numpy.fft.ifftshift", "jax._src.numpy.fft.ihfft", "jax._src.numpy.fft.irfft", "jax._src.numpy.fft.irfft2", "jax._src.numpy.fft.irfftn", "jax._src.numpy.fft.rfft", "jax._src.numpy.fft.rfft2", "jax._src.numpy.fft.rfftfreq", "jax._src.numpy.fft.rfftn" ], "n_typable": 86, "n_typed": 86, "n_any": 0, "n_untyped": 0, "n_functions": 18, "n_function_overloads": 18, "n_function_params": 68, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/numpy/linalg.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.numpy.linalg.matrix_rank", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.numpy.linalg.det", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.numpy.linalg.cond", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.numpy.linalg.matrix_transpose", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.numpy.linalg.vecdot", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.numpy.linalg.tensorsolve", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.numpy.linalg.diagonal", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.numpy.linalg.trace", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.numpy.linalg.eig", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.numpy.linalg.matmul", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.linalg.pinv", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.linalg.outer", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.numpy.linalg.tensorinv", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.numpy.linalg.eigh", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.numpy.linalg.inv", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.numpy.linalg.qr", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 3, "n_typable": 3, "n_params": 2, "n_function_overloads": 3, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.numpy.linalg.matrix_norm", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.numpy.linalg.eigvalsh", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.numpy.linalg.multi_dot", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.numpy.linalg.svdvals", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.numpy.linalg.slogdet", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.numpy.linalg.vector_norm", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.linalg.cholesky", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.numpy.linalg.svd", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 5, "n_typable": 6, "n_params": 5, "n_function_overloads": 5, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.numpy.linalg.solve", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.numpy.linalg.lstsq", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.numpy.linalg.matrix_power", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.numpy.linalg.eigvals", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.numpy.linalg.tensordot", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.numpy.linalg.cross", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.numpy.linalg.norm", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax._src.numpy.linalg", "names": [ "jax._src.numpy.linalg.cholesky", "jax._src.numpy.linalg.cond", "jax._src.numpy.linalg.cross", "jax._src.numpy.linalg.det", "jax._src.numpy.linalg.diagonal", "jax._src.numpy.linalg.eig", "jax._src.numpy.linalg.eigh", "jax._src.numpy.linalg.eigvals", "jax._src.numpy.linalg.eigvalsh", "jax._src.numpy.linalg.inv", "jax._src.numpy.linalg.lstsq", "jax._src.numpy.linalg.matmul", "jax._src.numpy.linalg.matrix_norm", "jax._src.numpy.linalg.matrix_power", "jax._src.numpy.linalg.matrix_rank", "jax._src.numpy.linalg.matrix_transpose", "jax._src.numpy.linalg.multi_dot", "jax._src.numpy.linalg.norm", "jax._src.numpy.linalg.outer", "jax._src.numpy.linalg.pinv", "jax._src.numpy.linalg.qr", "jax._src.numpy.linalg.slogdet", "jax._src.numpy.linalg.solve", "jax._src.numpy.linalg.svd", "jax._src.numpy.linalg.svdvals", "jax._src.numpy.linalg.tensordot", "jax._src.numpy.linalg.tensorinv", "jax._src.numpy.linalg.tensorsolve", "jax._src.numpy.linalg.trace", "jax._src.numpy.linalg.vecdot", "jax._src.numpy.linalg.vector_norm" ], "n_typable": 115, "n_typed": 111, "n_any": 0, "n_untyped": 4, "n_functions": 31, "n_function_overloads": 37, "n_function_params": 84, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/ops/scatter.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.ops.scatter.segment_sum", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.ops.scatter.segment_prod", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.ops.scatter.segment_max", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.ops.scatter.segment_min", "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 } ], "type_ignores": [], "name": "jax._src.ops.scatter", "names": [ "jax._src.ops.scatter.segment_max", "jax._src.ops.scatter.segment_min", "jax._src.ops.scatter.segment_prod", "jax._src.ops.scatter.segment_sum" ], "n_typable": 32, "n_typed": 32, "n_any": 0, "n_untyped": 0, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 28, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/profiler.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.profiler.stop_server", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.profiler.stop_trace", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "jax._src.profiler.device_memory_profile", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.profiler.start_trace", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "class", "name": "jax._src.profiler.TraceAnnotation", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax._src.profiler.annotate_function", "n_typed": 2, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "class", "name": "jax._src.profiler.StepTraceAnnotation", "methods": [ { "kind": "function", "name": "StepTraceAnnotation.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [], "attrs": [], "n_typable": 3, "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 2, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.profiler.ProfileOptions", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "function", "name": "jax._src.profiler.start_server", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.profiler.save_device_memory_profile", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.profiler.trace", "n_typed": 2, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax._src.profiler", "names": [ "jax._src.profiler.ProfileOptions", "jax._src.profiler.StepTraceAnnotation", "jax._src.profiler.TraceAnnotation", "jax._src.profiler.annotate_function", "jax._src.profiler.device_memory_profile", "jax._src.profiler.save_device_memory_profile", "jax._src.profiler.start_server", "jax._src.profiler.start_trace", "jax._src.profiler.stop_server", "jax._src.profiler.stop_trace", "jax._src.profiler.trace" ], "n_typable": 26, "n_typed": 16, "n_any": 0, "n_untyped": 10, "n_functions": 8, "n_function_overloads": 8, "n_function_params": 15, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 2, "n_classes": 3, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/profiler.py", "symbol_reports": [ { "kind": "attr", "name": "jax.profiler.ProfileEvent", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.profiler.ProfileData", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax.profiler.ProfilePlane", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [], "name": "jax.profiler", "names": [ "jax.profiler.ProfileData", "jax.profiler.ProfileEvent", "jax.profiler.ProfilePlane" ], "n_typable": 3, "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 3, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/cluster/vq.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.cluster.vq.vq", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.scipy.cluster.vq", "names": [ "jax._src.scipy.cluster.vq.vq" ], "n_typable": 4, "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 3, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/fft.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.fft.dct", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.fft.idct", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.fft.dctn", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.fft.idctn", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "type_ignores": [], "name": "jax._src.scipy.fft", "names": [ "jax._src.scipy.fft.dct", "jax._src.scipy.fft.dctn", "jax._src.scipy.fft.idct", "jax._src.scipy.fft.idctn" ], "n_typable": 24, "n_typed": 24, "n_any": 0, "n_untyped": 0, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 20, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/integrate.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.integrate.trapezoid", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax._src.scipy.integrate", "names": [ "jax._src.scipy.integrate.trapezoid" ], "n_typable": 5, "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 4, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/third_party/scipy/interpolate.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.third_party.scipy.interpolate.RegularGridInterpolator", "methods": [ { "kind": "function", "name": "RegularGridInterpolator.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "RegularGridInterpolator.__call__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "RegularGridInterpolator._evaluate_linear", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "RegularGridInterpolator._evaluate_nearest", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "RegularGridInterpolator._find_indices", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "RegularGridInterpolator.method", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "RegularGridInterpolator.bounds_error", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "RegularGridInterpolator.fill_value", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "RegularGridInterpolator.grid", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "RegularGridInterpolator.values", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 22, "n_typed": 0, "n_any": 0, "n_untyped": 22, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 12, "n_attrs": 5, "n_properties": 0 } ], "type_ignores": [], "name": "jax._src.third_party.scipy.interpolate", "names": [ "jax._src.third_party.scipy.interpolate.RegularGridInterpolator" ], "n_typable": 22, "n_typed": 0, "n_any": 0, "n_untyped": 22, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 12, "n_classes": 1, "n_attrs": 5, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/linalg.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.linalg.solve_triangular", "n_typed": 8, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax._src.scipy.linalg.det", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.linalg.expm", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.linalg.pascal", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.linalg.lu", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 3, "n_typable": 5, "n_params": 4, "n_function_overloads": 3, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.linalg.expm_frechet", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 3, "n_typable": 5, "n_params": 4, "n_function_overloads": 3, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.linalg.eigh", "n_typed": 11, "n_any": 0, "n_untyped": 0, "n_overloads": 4, "n_typable": 11, "n_params": 10, "n_function_overloads": 4, "n_function_params": 10 }, { "kind": "function", "name": "jax._src.scipy.linalg.inv", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.linalg.cho_solve", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.linalg.qr", "n_typed": 6, "n_any": 1, "n_untyped": 0, "n_overloads": 7, "n_typable": 7, "n_params": 6, "n_function_overloads": 7, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.scipy.linalg.cho_factor", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.linalg.rsf2csf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.linalg.toeplitz", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.linalg.schur", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.linalg.cholesky", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.linalg.hilbert", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.linalg.lu_solve", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.linalg.svd", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 4, "n_typable": 7, "n_params": 6, "n_function_overloads": 4, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.scipy.linalg.solve", "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 9, "n_params": 8, "n_function_overloads": 1, "n_function_params": 8 }, { "kind": "function", "name": "jax._src.scipy.linalg.lu_factor", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.linalg.eigh_tridiagonal", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.scipy.linalg.block_diag", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.linalg.sqrtm", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.linalg.solve_sylvester", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.linalg.hessenberg", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 2, "n_typable": 5, "n_params": 4, "n_function_overloads": 2, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.linalg.polar", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "type_ignores": [], "name": "jax._src.scipy.linalg", "names": [ "jax._src.scipy.linalg.block_diag", "jax._src.scipy.linalg.cho_factor", "jax._src.scipy.linalg.cho_solve", "jax._src.scipy.linalg.cholesky", "jax._src.scipy.linalg.det", "jax._src.scipy.linalg.eigh", "jax._src.scipy.linalg.eigh_tridiagonal", "jax._src.scipy.linalg.expm", "jax._src.scipy.linalg.expm_frechet", "jax._src.scipy.linalg.hessenberg", "jax._src.scipy.linalg.hilbert", "jax._src.scipy.linalg.inv", "jax._src.scipy.linalg.lu", "jax._src.scipy.linalg.lu_factor", "jax._src.scipy.linalg.lu_solve", "jax._src.scipy.linalg.pascal", "jax._src.scipy.linalg.polar", "jax._src.scipy.linalg.qr", "jax._src.scipy.linalg.rsf2csf", "jax._src.scipy.linalg.schur", "jax._src.scipy.linalg.solve", "jax._src.scipy.linalg.solve_sylvester", "jax._src.scipy.linalg.solve_triangular", "jax._src.scipy.linalg.sqrtm", "jax._src.scipy.linalg.svd", "jax._src.scipy.linalg.toeplitz" ], "n_typable": 134, "n_typed": 132, "n_any": 2, "n_untyped": 0, "n_functions": 26, "n_function_overloads": 43, "n_function_params": 108, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/third_party/scipy/linalg.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.third_party.scipy.linalg.funm", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.third_party.scipy.linalg", "names": [ "jax._src.third_party.scipy.linalg.funm" ], "n_typable": 4, "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 3, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/ndimage.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.ndimage.map_coordinates", "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "type_ignores": [], "name": "jax._src.scipy.ndimage", "names": [ "jax._src.scipy.ndimage.map_coordinates" ], "n_typable": 6, "n_typed": 5, "n_any": 0, "n_untyped": 1, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 5, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/optimize/minimize.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.optimize.minimize.minimize", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "class", "name": "jax._src.scipy.optimize.minimize.OptimizeResults", "methods": [], "properties": [], "attrs": [ { "kind": "attr", "name": "OptimizeResults.x", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "OptimizeResults.success", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "OptimizeResults.status", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "OptimizeResults.fun", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "OptimizeResults.jac", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "OptimizeResults.hess_inv", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "OptimizeResults.nfev", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "OptimizeResults.njev", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "OptimizeResults.nit", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 9, "n_properties": 0 } ], "type_ignores": [], "name": "jax._src.scipy.optimize.minimize", "names": [ "jax._src.scipy.optimize.minimize.OptimizeResults", "jax._src.scipy.optimize.minimize.minimize" ], "n_typable": 7, "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 6, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 1, "n_attrs": 9, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/signal.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.signal.fftconvolve", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.signal.convolve", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.signal.correlate", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.signal.correlate2d", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.scipy.signal.welch", "n_typed": 12, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 12, "n_params": 11, "n_function_overloads": 1, "n_function_params": 11 }, { "kind": "function", "name": "jax._src.scipy.signal.stft", "n_typed": 12, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 12, "n_params": 11, "n_function_overloads": 1, "n_function_params": 11 }, { "kind": "function", "name": "jax._src.scipy.signal.istft", "n_typed": 11, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 11, "n_params": 10, "n_function_overloads": 1, "n_function_params": 10 }, { "kind": "function", "name": "jax._src.scipy.signal.convolve2d", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.scipy.signal.detrend", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.signal.csd", "n_typed": 13, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 13, "n_params": 12, "n_function_overloads": 1, "n_function_params": 12 } ], "type_ignores": [], "name": "jax._src.scipy.signal", "names": [ "jax._src.scipy.signal.convolve", "jax._src.scipy.signal.convolve2d", "jax._src.scipy.signal.correlate", "jax._src.scipy.signal.correlate2d", "jax._src.scipy.signal.csd", "jax._src.scipy.signal.detrend", "jax._src.scipy.signal.fftconvolve", "jax._src.scipy.signal.istft", "jax._src.scipy.signal.stft", "jax._src.scipy.signal.welch" ], "n_typable": 85, "n_typed": 85, "n_any": 0, "n_untyped": 0, "n_functions": 10, "n_function_overloads": 10, "n_function_params": 75, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/sparse/linalg.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.sparse.linalg.gmres", "n_typed": 0, "n_any": 0, "n_untyped": 10, "n_overloads": 1, "n_typable": 10, "n_params": 9, "n_function_overloads": 1, "n_function_params": 9 }, { "kind": "function", "name": "jax._src.scipy.sparse.linalg.cg", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.scipy.sparse.linalg.bicgstab", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 } ], "type_ignores": [], "name": "jax._src.scipy.sparse.linalg", "names": [ "jax._src.scipy.sparse.linalg.bicgstab", "jax._src.scipy.sparse.linalg.cg", "jax._src.scipy.sparse.linalg.gmres" ], "n_typable": 26, "n_typed": 0, "n_any": 0, "n_untyped": 26, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 23, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/spatial/transform.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.scipy.spatial.transform.Slerp", "methods": [ { "kind": "function", "name": "Slerp.init", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Slerp.__call__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [], "attrs": [ { "kind": "attr", "name": "Slerp.times", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Slerp.timedelta", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Slerp.rotations", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "Slerp.rotvecs", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 5, "n_typed": 3, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 2, "n_method_overloads": 2, "n_method_params": 3, "n_attrs": 4, "n_properties": 0 }, { "kind": "class", "name": "jax._src.scipy.spatial.transform.Rotation", "methods": [ { "kind": "function", "name": "Rotation.concatenate", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Rotation.from_euler", "n_typed": 3, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "Rotation.from_matrix", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Rotation.from_mrp", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Rotation.from_quat", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Rotation.from_rotvec", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Rotation.identity", "n_typed": 1, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Rotation.random", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Rotation.__getitem__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Rotation.__len__", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Rotation.__mul__", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Rotation.apply", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Rotation.as_euler", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Rotation.as_matrix", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Rotation.as_mrp", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Rotation.as_rotvec", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Rotation.as_quat", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Rotation.inv", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Rotation.magnitude", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Rotation.mean", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "Rotation.single", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "Rotation.quat", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 44, "n_typed": 28, "n_any": 0, "n_untyped": 16, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 20, "n_method_overloads": 20, "n_method_params": 23, "n_attrs": 1, "n_properties": 1 } ], "type_ignores": [], "name": "jax._src.scipy.spatial.transform", "names": [ "jax._src.scipy.spatial.transform.Rotation", "jax._src.scipy.spatial.transform.Slerp" ], "n_typable": 49, "n_typed": 31, "n_any": 0, "n_untyped": 18, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 22, "n_method_overloads": 22, "n_method_params": 26, "n_classes": 2, "n_attrs": 5, "n_properties": 1, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/special.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.special.bessel_jn", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.special.log_ndtr", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.special.factorial", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.special.i1", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.multigammaln", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.special.ndtri", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.softmax", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.special.i0e", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.expn", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.special.exp1", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.beta", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.special.gammaincc", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.special.log_softmax", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.special.gammasgn", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.i1e", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.poch", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.special.hyp2f1", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.special.betainc", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.special.digamma", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.expit", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.sph_harm_y", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "jax._src.scipy.special.gammainc", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.special.sici", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.betaln", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.special.bernoulli", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.xlogy", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.special.hyp1f1", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.special.entr", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.logit", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.xlog1py", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.special.gamma", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.zeta", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.special.gammaln", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.erf", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.spence", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.expi", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.erfc", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.ndtr", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.erfinv", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.polygamma", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.special.i0", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.scipy.special.rel_entr", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.special.kl_div", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [], "name": "jax._src.scipy.special", "names": [ "jax._src.scipy.special.bernoulli", "jax._src.scipy.special.bessel_jn", "jax._src.scipy.special.beta", "jax._src.scipy.special.betainc", "jax._src.scipy.special.betaln", "jax._src.scipy.special.digamma", "jax._src.scipy.special.entr", "jax._src.scipy.special.erf", "jax._src.scipy.special.erfc", "jax._src.scipy.special.erfinv", "jax._src.scipy.special.exp1", "jax._src.scipy.special.expi", "jax._src.scipy.special.expit", "jax._src.scipy.special.expn", "jax._src.scipy.special.factorial", "jax._src.scipy.special.gamma", "jax._src.scipy.special.gammainc", "jax._src.scipy.special.gammaincc", "jax._src.scipy.special.gammaln", "jax._src.scipy.special.gammasgn", "jax._src.scipy.special.hyp1f1", "jax._src.scipy.special.hyp2f1", "jax._src.scipy.special.i0", "jax._src.scipy.special.i0e", "jax._src.scipy.special.i1", "jax._src.scipy.special.i1e", "jax._src.scipy.special.kl_div", "jax._src.scipy.special.log_ndtr", "jax._src.scipy.special.log_softmax", "jax._src.scipy.special.logit", "jax._src.scipy.special.multigammaln", "jax._src.scipy.special.ndtr", "jax._src.scipy.special.ndtri", "jax._src.scipy.special.poch", "jax._src.scipy.special.polygamma", "jax._src.scipy.special.rel_entr", "jax._src.scipy.special.sici", "jax._src.scipy.special.softmax", "jax._src.scipy.special.spence", "jax._src.scipy.special.sph_harm_y", "jax._src.scipy.special.xlog1py", "jax._src.scipy.special.xlogy", "jax._src.scipy.special.zeta" ], "n_typable": 117, "n_typed": 117, "n_any": 0, "n_untyped": 0, "n_functions": 43, "n_function_overloads": 43, "n_function_params": 74, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/ops/special.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.ops.special.logsumexp", "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_overloads": 3, "n_typable": 7, "n_params": 6, "n_function_overloads": 3, "n_function_params": 6 } ], "type_ignores": [], "name": "jax._src.ops.special", "names": [ "jax._src.ops.special.logsumexp" ], "n_typable": 7, "n_typed": 7, "n_any": 0, "n_untyped": 0, "n_functions": 1, "n_function_overloads": 3, "n_function_params": 6, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/third_party/scipy/special.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.third_party.scipy.special.fresnel", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "type_ignores": [], "name": "jax._src.third_party.scipy.special", "names": [ "jax._src.third_party.scipy.special.fresnel" ], "n_typable": 2, "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_functions": 1, "n_function_overloads": 1, "n_function_params": 1, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/_core.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats._core.mode", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats._core.sem", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.stats._core.rankdata", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [ { "kind": "type", "rules": [ "bad-override", "assignment" ] } ], "name": "jax._src.scipy.stats._core", "names": [ "jax._src.scipy.stats._core.mode", "jax._src.scipy.stats._core.rankdata", "jax._src.scipy.stats._core.sem" ], "n_typable": 16, "n_typed": 16, "n_any": 0, "n_untyped": 0, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 13, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 1 }, { "path": "jax/_src/scipy/stats/kde.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.scipy.stats.kde.gaussian_kde", "methods": [ { "kind": "function", "name": "gaussian_kde.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "gaussian_kde._setattr", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "gaussian_kde.tree_flatten", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "gaussian_kde.tree_unflatten", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "gaussian_kde.evaluate", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "gaussian_kde.__call__", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "gaussian_kde.integrate_gaussian", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "gaussian_kde.integrate_box_1d", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "gaussian_kde.integrate_kde", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "gaussian_kde.resample", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "gaussian_kde.pdf", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "gaussian_kde.logpdf", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "gaussian_kde.integrate_box", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "gaussian_kde.set_bandwidth", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "gaussian_kde._reshape_points", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 } ], "properties": [ { "kind": "property", "name": "gaussian_kde.d", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "gaussian_kde.n", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "gaussian_kde.neff", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "gaussian_kde.dataset", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "gaussian_kde.weights", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "gaussian_kde.covariance", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "gaussian_kde.inv_cov", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 40, "n_typed": 1, "n_any": 0, "n_untyped": 39, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 15, "n_method_overloads": 15, "n_method_params": 23, "n_attrs": 5, "n_properties": 2 } ], "type_ignores": [], "name": "jax._src.scipy.stats.kde", "names": [ "jax._src.scipy.stats.kde.gaussian_kde" ], "n_typable": 40, "n_typed": 1, "n_any": 0, "n_untyped": 39, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 15, "n_method_overloads": 15, "n_method_params": 23, "n_classes": 1, "n_attrs": 5, "n_properties": 2, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/bernoulli.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.bernoulli.logpmf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.bernoulli.pmf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.bernoulli.cdf", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.stats.bernoulli.ppf", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [], "name": "jax._src.scipy.stats.bernoulli", "names": [ "jax._src.scipy.stats.bernoulli.cdf", "jax._src.scipy.stats.bernoulli.logpmf", "jax._src.scipy.stats.bernoulli.pmf", "jax._src.scipy.stats.bernoulli.ppf" ], "n_typable": 14, "n_typed": 14, "n_any": 0, "n_untyped": 0, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 10, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/beta.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.beta.cdf", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.stats.beta.logsf", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.stats.beta.logcdf", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.stats.beta.sf", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.stats.beta.logpdf", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.stats.beta.pdf", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "type_ignores": [], "name": "jax._src.scipy.stats.beta", "names": [ "jax._src.scipy.stats.beta.cdf", "jax._src.scipy.stats.beta.logcdf", "jax._src.scipy.stats.beta.logpdf", "jax._src.scipy.stats.beta.logsf", "jax._src.scipy.stats.beta.pdf", "jax._src.scipy.stats.beta.sf" ], "n_typable": 36, "n_typed": 36, "n_any": 0, "n_untyped": 0, "n_functions": 6, "n_function_overloads": 6, "n_function_params": 30, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/betabinom.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.betabinom.logpmf", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.stats.betabinom.pmf", "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "type_ignores": [], "name": "jax._src.scipy.stats.betabinom", "names": [ "jax._src.scipy.stats.betabinom.logpmf", "jax._src.scipy.stats.betabinom.pmf" ], "n_typable": 12, "n_typed": 12, "n_any": 0, "n_untyped": 0, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 10, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/binom.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.binom.logpmf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.binom.pmf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax._src.scipy.stats.binom", "names": [ "jax._src.scipy.stats.binom.logpmf", "jax._src.scipy.stats.binom.pmf" ], "n_typable": 10, "n_typed": 10, "n_any": 0, "n_untyped": 0, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 8, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/cauchy.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.cauchy.logpdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.cauchy.cdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.cauchy.logcdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.cauchy.sf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.cauchy.ppf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.cauchy.isf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.cauchy.logsf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.cauchy.pdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.scipy.stats.cauchy", "names": [ "jax._src.scipy.stats.cauchy.cdf", "jax._src.scipy.stats.cauchy.isf", "jax._src.scipy.stats.cauchy.logcdf", "jax._src.scipy.stats.cauchy.logpdf", "jax._src.scipy.stats.cauchy.logsf", "jax._src.scipy.stats.cauchy.pdf", "jax._src.scipy.stats.cauchy.ppf", "jax._src.scipy.stats.cauchy.sf" ], "n_typable": 32, "n_typed": 32, "n_any": 0, "n_untyped": 0, "n_functions": 8, "n_function_overloads": 8, "n_function_params": 24, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/chi2.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.chi2.cdf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.chi2.logsf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.chi2.logcdf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.chi2.sf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.chi2.logpdf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.chi2.pdf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax._src.scipy.stats.chi2", "names": [ "jax._src.scipy.stats.chi2.cdf", "jax._src.scipy.stats.chi2.logcdf", "jax._src.scipy.stats.chi2.logpdf", "jax._src.scipy.stats.chi2.logsf", "jax._src.scipy.stats.chi2.pdf", "jax._src.scipy.stats.chi2.sf" ], "n_typable": 30, "n_typed": 30, "n_any": 0, "n_untyped": 0, "n_functions": 6, "n_function_overloads": 6, "n_function_params": 24, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/dirichlet.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.dirichlet.logpdf", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.stats.dirichlet.pdf", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [], "name": "jax._src.scipy.stats.dirichlet", "names": [ "jax._src.scipy.stats.dirichlet.logpdf", "jax._src.scipy.stats.dirichlet.pdf" ], "n_typable": 6, "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 4, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/expon.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.expon.cdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.expon.logsf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.expon.logcdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.expon.sf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.expon.ppf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.expon.logpdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.expon.pdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.scipy.stats.expon", "names": [ "jax._src.scipy.stats.expon.cdf", "jax._src.scipy.stats.expon.logcdf", "jax._src.scipy.stats.expon.logpdf", "jax._src.scipy.stats.expon.logsf", "jax._src.scipy.stats.expon.pdf", "jax._src.scipy.stats.expon.ppf", "jax._src.scipy.stats.expon.sf" ], "n_typable": 28, "n_typed": 28, "n_any": 0, "n_untyped": 0, "n_functions": 7, "n_function_overloads": 7, "n_function_params": 21, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/gamma.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.gamma.cdf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.gamma.logsf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.gamma.logcdf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.gamma.sf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.gamma.logpdf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.gamma.pdf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax._src.scipy.stats.gamma", "names": [ "jax._src.scipy.stats.gamma.cdf", "jax._src.scipy.stats.gamma.logcdf", "jax._src.scipy.stats.gamma.logpdf", "jax._src.scipy.stats.gamma.logsf", "jax._src.scipy.stats.gamma.pdf", "jax._src.scipy.stats.gamma.sf" ], "n_typable": 30, "n_typed": 30, "n_any": 0, "n_untyped": 0, "n_functions": 6, "n_function_overloads": 6, "n_function_params": 24, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/gennorm.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.gennorm.cdf", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.stats.gennorm.logpdf", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.stats.gennorm.pdf", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [], "name": "jax._src.scipy.stats.gennorm", "names": [ "jax._src.scipy.stats.gennorm.cdf", "jax._src.scipy.stats.gennorm.logpdf", "jax._src.scipy.stats.gennorm.pdf" ], "n_typable": 9, "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 6, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/geom.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.geom.logpmf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.geom.pmf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.scipy.stats.geom", "names": [ "jax._src.scipy.stats.geom.logpmf", "jax._src.scipy.stats.geom.pmf" ], "n_typable": 8, "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 6, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/gumbel_l.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.gumbel_l.cdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.gumbel_l.logsf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.gumbel_l.logcdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.gumbel_l.sf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.gumbel_l.ppf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.gumbel_l.logpdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.gumbel_l.pdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.scipy.stats.gumbel_l", "names": [ "jax._src.scipy.stats.gumbel_l.cdf", "jax._src.scipy.stats.gumbel_l.logcdf", "jax._src.scipy.stats.gumbel_l.logpdf", "jax._src.scipy.stats.gumbel_l.logsf", "jax._src.scipy.stats.gumbel_l.pdf", "jax._src.scipy.stats.gumbel_l.ppf", "jax._src.scipy.stats.gumbel_l.sf" ], "n_typable": 28, "n_typed": 28, "n_any": 0, "n_untyped": 0, "n_functions": 7, "n_function_overloads": 7, "n_function_params": 21, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/gumbel_r.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.gumbel_r.cdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.gumbel_r.logsf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.gumbel_r.logcdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.gumbel_r.sf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.gumbel_r.ppf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.gumbel_r.logpdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.gumbel_r.pdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.scipy.stats.gumbel_r", "names": [ "jax._src.scipy.stats.gumbel_r.cdf", "jax._src.scipy.stats.gumbel_r.logcdf", "jax._src.scipy.stats.gumbel_r.logpdf", "jax._src.scipy.stats.gumbel_r.logsf", "jax._src.scipy.stats.gumbel_r.pdf", "jax._src.scipy.stats.gumbel_r.ppf", "jax._src.scipy.stats.gumbel_r.sf" ], "n_typable": 28, "n_typed": 28, "n_any": 0, "n_untyped": 0, "n_functions": 7, "n_function_overloads": 7, "n_function_params": 21, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/laplace.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.laplace.cdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.laplace.logpdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.laplace.pdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.scipy.stats.laplace", "names": [ "jax._src.scipy.stats.laplace.cdf", "jax._src.scipy.stats.laplace.logpdf", "jax._src.scipy.stats.laplace.pdf" ], "n_typable": 12, "n_typed": 12, "n_any": 0, "n_untyped": 0, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 9, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/logistic.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.logistic.cdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.logistic.sf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.logistic.ppf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.logistic.isf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.logistic.logpdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.logistic.pdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.scipy.stats.logistic", "names": [ "jax._src.scipy.stats.logistic.cdf", "jax._src.scipy.stats.logistic.isf", "jax._src.scipy.stats.logistic.logpdf", "jax._src.scipy.stats.logistic.pdf", "jax._src.scipy.stats.logistic.ppf", "jax._src.scipy.stats.logistic.sf" ], "n_typable": 24, "n_typed": 24, "n_any": 0, "n_untyped": 0, "n_functions": 6, "n_function_overloads": 6, "n_function_params": 18, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/multinomial.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.multinomial.logpmf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.multinomial.pmf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.scipy.stats.multinomial", "names": [ "jax._src.scipy.stats.multinomial.logpmf", "jax._src.scipy.stats.multinomial.pmf" ], "n_typable": 8, "n_typed": 8, "n_any": 0, "n_untyped": 0, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 6, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/multivariate_normal.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.multivariate_normal.logpdf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.multivariate_normal.pdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.scipy.stats.multivariate_normal", "names": [ "jax._src.scipy.stats.multivariate_normal.logpdf", "jax._src.scipy.stats.multivariate_normal.pdf" ], "n_typable": 9, "n_typed": 9, "n_any": 0, "n_untyped": 0, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 7, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/nbinom.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.nbinom.logpmf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.nbinom.pmf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax._src.scipy.stats.nbinom", "names": [ "jax._src.scipy.stats.nbinom.logpmf", "jax._src.scipy.stats.nbinom.pmf" ], "n_typable": 10, "n_typed": 10, "n_any": 0, "n_untyped": 0, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 8, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/norm.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.norm.logpdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.norm.cdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.norm.logcdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.norm.sf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.norm.ppf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.norm.isf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.norm.logsf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.norm.pdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.scipy.stats.norm", "names": [ "jax._src.scipy.stats.norm.cdf", "jax._src.scipy.stats.norm.isf", "jax._src.scipy.stats.norm.logcdf", "jax._src.scipy.stats.norm.logpdf", "jax._src.scipy.stats.norm.logsf", "jax._src.scipy.stats.norm.pdf", "jax._src.scipy.stats.norm.ppf", "jax._src.scipy.stats.norm.sf" ], "n_typable": 32, "n_typed": 32, "n_any": 0, "n_untyped": 0, "n_functions": 8, "n_function_overloads": 8, "n_function_params": 24, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/pareto.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.pareto.cdf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.pareto.logsf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.pareto.logcdf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.pareto.sf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.pareto.ppf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.pareto.logpdf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.pareto.pdf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax._src.scipy.stats.pareto", "names": [ "jax._src.scipy.stats.pareto.cdf", "jax._src.scipy.stats.pareto.logcdf", "jax._src.scipy.stats.pareto.logpdf", "jax._src.scipy.stats.pareto.logsf", "jax._src.scipy.stats.pareto.pdf", "jax._src.scipy.stats.pareto.ppf", "jax._src.scipy.stats.pareto.sf" ], "n_typable": 35, "n_typed": 35, "n_any": 0, "n_untyped": 0, "n_functions": 7, "n_function_overloads": 7, "n_function_params": 28, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/poisson.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.poisson.entropy", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.stats.poisson.logpmf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.poisson.pmf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.poisson.cdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.scipy.stats.poisson", "names": [ "jax._src.scipy.stats.poisson.cdf", "jax._src.scipy.stats.poisson.entropy", "jax._src.scipy.stats.poisson.logpmf", "jax._src.scipy.stats.poisson.pmf" ], "n_typable": 15, "n_typed": 15, "n_any": 0, "n_untyped": 0, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 11, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/t.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.t.logpdf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.scipy.stats.t.pdf", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 } ], "type_ignores": [], "name": "jax._src.scipy.stats.t", "names": [ "jax._src.scipy.stats.t.logpdf", "jax._src.scipy.stats.t.pdf" ], "n_typable": 10, "n_typed": 10, "n_any": 0, "n_untyped": 0, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 8, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/truncnorm.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.truncnorm.cdf", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.stats.truncnorm.logsf", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.stats.truncnorm.logcdf", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.stats.truncnorm.sf", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.stats.truncnorm.logpdf", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.scipy.stats.truncnorm.pdf", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 } ], "type_ignores": [], "name": "jax._src.scipy.stats.truncnorm", "names": [ "jax._src.scipy.stats.truncnorm.cdf", "jax._src.scipy.stats.truncnorm.logcdf", "jax._src.scipy.stats.truncnorm.logpdf", "jax._src.scipy.stats.truncnorm.logsf", "jax._src.scipy.stats.truncnorm.pdf", "jax._src.scipy.stats.truncnorm.sf" ], "n_typable": 36, "n_typed": 0, "n_any": 0, "n_untyped": 36, "n_functions": 6, "n_function_overloads": 6, "n_function_params": 30, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/uniform.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.uniform.logpdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.uniform.cdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.uniform.ppf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.scipy.stats.uniform.pdf", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.scipy.stats.uniform", "names": [ "jax._src.scipy.stats.uniform.cdf", "jax._src.scipy.stats.uniform.logpdf", "jax._src.scipy.stats.uniform.pdf", "jax._src.scipy.stats.uniform.ppf" ], "n_typable": 16, "n_typed": 16, "n_any": 0, "n_untyped": 0, "n_functions": 4, "n_function_overloads": 4, "n_function_params": 12, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/vonmises.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.vonmises.logpdf", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.stats.vonmises.pdf", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [], "name": "jax._src.scipy.stats.vonmises", "names": [ "jax._src.scipy.stats.vonmises.logpdf", "jax._src.scipy.stats.vonmises.pdf" ], "n_typable": 6, "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 4, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/scipy/stats/wrapcauchy.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.scipy.stats.wrapcauchy.logpdf", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.scipy.stats.wrapcauchy.pdf", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "type_ignores": [], "name": "jax._src.scipy.stats.wrapcauchy", "names": [ "jax._src.scipy.stats.wrapcauchy.logpdf", "jax._src.scipy.stats.wrapcauchy.pdf" ], "n_typable": 6, "n_typed": 6, "n_any": 0, "n_untyped": 0, "n_functions": 2, "n_function_overloads": 2, "n_function_params": 4, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/sharding.py", "symbol_reports": [ { "kind": "attr", "name": "jax.sharding.PartitionSpec", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "type_ignores": [], "name": "jax.sharding", "names": [ "jax.sharding.PartitionSpec" ], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/sharding.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.sharding.Sharding", "methods": [ { "kind": "function", "name": "Sharding.with_memory_kind", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Sharding._to_xla_hlo_sharding", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Sharding._to_sdy_sharding", "n_typed": 1, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Sharding.addressable_devices_indices_map", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Sharding.devices_indices_map", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Sharding.shard_shape", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Sharding.is_equivalent_to", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [ { "kind": "property", "name": "Sharding.device_set", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Sharding.is_fully_replicated", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Sharding.is_fully_addressable", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Sharding.num_devices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Sharding.memory_kind", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Sharding._device_assignment", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Sharding._internal_device_list", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Sharding._is_concrete", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Sharding.addressable_devices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Sharding.has_addressable_devices", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "property", "name": "Sharding._addressable_device_assignment", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [], "n_typable": 26, "n_typed": 25, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 8, "n_attrs": 0, "n_properties": 11 } ], "type_ignores": [], "name": "jax._src.sharding", "names": [ "jax._src.sharding.Sharding" ], "n_typable": 26, "n_typed": 25, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 7, "n_method_overloads": 7, "n_method_params": 8, "n_classes": 1, "n_attrs": 0, "n_properties": 11, "n_type_ignores": 0 }, { "path": "jax/_src/stages.py", "symbol_reports": [ { "kind": "class", "name": "jax._src.stages.Lowered", "methods": [ { "kind": "function", "name": "Lowered.__init__", "n_typed": 3, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 7, "n_params": 6, "n_function_overloads": 1, "n_function_params": 6 }, { "kind": "function", "name": "Lowered.compile", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Lowered.as_text", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Lowered.compiler_ir", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "Lowered.cost_analysis", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 } ], "properties": [ { "kind": "property", "name": "Lowered.in_avals", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Lowered.out_info", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "Lowered._lowering", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Lowered.args_info", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Lowered.out_tree", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Lowered._no_kwargs", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Lowered._in_types", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Lowered._out_types", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 24, "n_typed": 17, "n_any": 1, "n_untyped": 6, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 5, "n_method_overloads": 5, "n_method_params": 11, "n_attrs": 6, "n_properties": 2 }, { "kind": "class", "name": "jax._src.stages.Wrapped", "methods": [], "properties": [], "attrs": [], "n_typable": 0, "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 0, "n_properties": 0 }, { "kind": "class", "name": "jax._src.stages.Compiled", "methods": [ { "kind": "function", "name": "Compiled.__init__", "n_typed": 1, "n_any": 0, "n_untyped": 7, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "Compiled.as_text", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Compiled.cost_analysis", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Compiled.memory_analysis", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Compiled.runtime_executable", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Compiled._input_shardings_flat", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Compiled._input_layouts_flat", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 1, "n_params": 0, "n_function_overloads": 1, "n_function_params": 0 }, { "kind": "function", "name": "Compiled.call", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Compiled.__call__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [ { "kind": "property", "name": "Compiled.in_avals", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Compiled.out_info", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Compiled.input_shardings", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Compiled.output_shardings", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Compiled.input_formats", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Compiled._output_formats_flat", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Compiled.output_formats", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "Compiled.args_info", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Compiled.out_tree", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Compiled._executable", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Compiled._no_kwargs", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "Compiled._params", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "n_typable": 32, "n_typed": 9, "n_any": 1, "n_untyped": 22, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 9, "n_method_overloads": 9, "n_method_params": 11, "n_attrs": 5, "n_properties": 7 }, { "kind": "class", "name": "jax._src.stages.ArgInfo", "methods": [], "properties": [ { "kind": "property", "name": "ArgInfo.shape", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "ArgInfo.dtype", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "ArgInfo._aval", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "ArgInfo.donated", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 } ], "n_typable": 2, "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_attrs": 2, "n_properties": 2 }, { "kind": "attr", "name": "jax._src.stages.CompilerOptions", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "class", "name": "jax._src.stages.Traced", "methods": [ { "kind": "function", "name": "Traced.__init__", "n_typed": 0, "n_any": 0, "n_untyped": 6, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "Traced.__call__", "n_typed": 0, "n_any": 0, "n_untyped": 3, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "Traced.lower", "n_typed": 2, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 } ], "properties": [ { "kind": "property", "name": "Traced.out_avals", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "property", "name": "Traced.lojax", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "attrs": [ { "kind": "attr", "name": "Traced.out_tree", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Traced.jaxpr", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Traced.fun_name", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Traced.args_info", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Traced.out_info", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "Traced._num_consts", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "n_typable": 20, "n_typed": 3, "n_any": 0, "n_untyped": 17, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 3, "n_method_overloads": 3, "n_method_params": 9, "n_attrs": 6, "n_properties": 2 } ], "type_ignores": [], "name": "jax._src.stages", "names": [ "jax._src.stages.ArgInfo", "jax._src.stages.Compiled", "jax._src.stages.CompilerOptions", "jax._src.stages.Lowered", "jax._src.stages.Traced", "jax._src.stages.Wrapped" ], "n_typable": 78, "n_typed": 29, "n_any": 2, "n_untyped": 47, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 17, "n_method_overloads": 17, "n_method_params": 31, "n_classes": 5, "n_attrs": 20, "n_properties": 13, "n_type_ignores": 0 }, { "path": "jax/_src/public_test_util.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.public_test_util.check_grads", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.public_test_util.check_vjp", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 }, { "kind": "function", "name": "jax._src.public_test_util.check_jvp", "n_typed": 0, "n_any": 0, "n_untyped": 8, "n_overloads": 1, "n_typable": 8, "n_params": 7, "n_function_overloads": 1, "n_function_params": 7 } ], "type_ignores": [], "name": "jax._src.public_test_util", "names": [ "jax._src.public_test_util.check_grads", "jax._src.public_test_util.check_jvp", "jax._src.public_test_util.check_vjp" ], "n_typable": 24, "n_typed": 0, "n_any": 0, "n_untyped": 24, "n_functions": 3, "n_function_overloads": 3, "n_function_params": 21, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/tree.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.tree.static", "n_typed": 0, "n_any": 0, "n_untyped": 2, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.tree.reduce", "n_typed": 4, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.tree.reduce_associative", "n_typed": 4, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.tree.map_with_path", "n_typed": 3, "n_any": 3, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "function", "name": "jax._src.tree.flatten_with_path", "n_typed": 3, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.tree.leaves_with_path", "n_typed": 3, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.tree.map", "n_typed": 2, "n_any": 3, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.tree.unflatten", "n_typed": 2, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.tree.structure", "n_typed": 2, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.tree.flatten", "n_typed": 2, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.tree.all", "n_typed": 2, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.tree.broadcast", "n_typed": 1, "n_any": 3, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.tree.leaves", "n_typed": 2, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.tree.transpose", "n_typed": 2, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [], "name": "jax._src.tree", "names": [ "jax._src.tree.all", "jax._src.tree.broadcast", "jax._src.tree.flatten", "jax._src.tree.flatten_with_path", "jax._src.tree.leaves", "jax._src.tree.leaves_with_path", "jax._src.tree.map", "jax._src.tree.map_with_path", "jax._src.tree.reduce", "jax._src.tree.reduce_associative", "jax._src.tree.static", "jax._src.tree.structure", "jax._src.tree.transpose", "jax._src.tree.unflatten" ], "n_typable": 54, "n_typed": 32, "n_any": 20, "n_untyped": 2, "n_functions": 14, "n_function_overloads": 14, "n_function_params": 40, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 0, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/tree_util.py", "symbol_reports": [ { "kind": "function", "name": "jax._src.tree_util.tree_unflatten", "n_typed": 2, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "class", "name": "jax._src.tree_util.Partial", "methods": [ { "kind": "function", "name": "Partial.__new__", "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "properties": [], "attrs": [], "n_typable": 4, "n_typed": 0, "n_any": 0, "n_untyped": 4, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 3, "n_attrs": 0, "n_properties": 0 }, { "kind": "attr", "name": "jax._src.tree_util.default_registry", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.tree_util.KeyPath", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "attr", "name": "jax._src.tree_util.KeyEntry", "n_typed": 0, "n_any": 0, "n_untyped": 0, "n_typable": 0 }, { "kind": "function", "name": "jax._src.tree_util.tree_reduce_associative", "n_typed": 4, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.tree_util.tree_all", "n_typed": 2, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.tree_util.tree_flatten", "n_typed": 2, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.tree_util.tree_reduce", "n_typed": 4, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "attr", "name": "jax._src.tree_util.DictKey", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "attr", "name": "jax._src.tree_util.GetAttrKey", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax._src.tree_util.all_leaves", "n_typed": 3, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.tree_util.tree_map_with_path", "n_typed": 3, "n_any": 3, "n_untyped": 0, "n_overloads": 1, "n_typable": 6, "n_params": 5, "n_function_overloads": 1, "n_function_params": 5 }, { "kind": "attr", "name": "jax._src.tree_util.SequenceKey", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax._src.tree_util.tree_broadcast", "n_typed": 1, "n_any": 3, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.tree_util.register_dataclass", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.tree_util.treedef_is_leaf", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "attr", "name": "jax._src.tree_util.FlattenedIndexKey", "n_typed": 0, "n_any": 1, "n_untyped": 0, "n_typable": 1 }, { "kind": "function", "name": "jax._src.tree_util.tree_map", "n_typed": 2, "n_any": 3, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.tree_util.treedef_children", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.tree_util.treedef_tuple", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.tree_util.register_pytree_node", "n_typed": 5, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.tree_util.register_pytree_with_keys", "n_typed": 4, "n_any": 0, "n_untyped": 1, "n_overloads": 1, "n_typable": 5, "n_params": 4, "n_function_overloads": 1, "n_function_params": 4 }, { "kind": "function", "name": "jax._src.tree_util.register_static", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.tree_util.keystr", "n_typed": 4, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.tree_util.register_pytree_with_keys_class", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.tree_util.register_pytree_node_class", "n_typed": 2, "n_any": 0, "n_untyped": 0, "n_overloads": 1, "n_typable": 2, "n_params": 1, "n_function_overloads": 1, "n_function_params": 1 }, { "kind": "function", "name": "jax._src.tree_util.tree_leaves_with_path", "n_typed": 3, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.tree_util.tree_leaves", "n_typed": 2, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.tree_util.tree_structure", "n_typed": 2, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 3, "n_params": 2, "n_function_overloads": 1, "n_function_params": 2 }, { "kind": "function", "name": "jax._src.tree_util.tree_flatten_with_path", "n_typed": 3, "n_any": 1, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 }, { "kind": "function", "name": "jax._src.tree_util.tree_transpose", "n_typed": 2, "n_any": 2, "n_untyped": 0, "n_overloads": 1, "n_typable": 4, "n_params": 3, "n_function_overloads": 1, "n_function_params": 3 } ], "type_ignores": [ { "kind": "pyrefly", "rules": [ "missing-attribute" ] }, { "kind": "pyrefly", "rules": [ "missing-attribute" ] }, { "kind": "pyrefly", "rules": [ "missing-attribute" ] } ], "name": "jax._src.tree_util", "names": [ "jax._src.tree_util.DictKey", "jax._src.tree_util.FlattenedIndexKey", "jax._src.tree_util.GetAttrKey", "jax._src.tree_util.KeyEntry", "jax._src.tree_util.KeyPath", "jax._src.tree_util.Partial", "jax._src.tree_util.SequenceKey", "jax._src.tree_util.all_leaves", "jax._src.tree_util.default_registry", "jax._src.tree_util.keystr", "jax._src.tree_util.register_dataclass", "jax._src.tree_util.register_pytree_node", "jax._src.tree_util.register_pytree_node_class", "jax._src.tree_util.register_pytree_with_keys", "jax._src.tree_util.register_pytree_with_keys_class", "jax._src.tree_util.register_static", "jax._src.tree_util.tree_all", "jax._src.tree_util.tree_broadcast", "jax._src.tree_util.tree_flatten", "jax._src.tree_util.tree_flatten_with_path", "jax._src.tree_util.tree_leaves", "jax._src.tree_util.tree_leaves_with_path", "jax._src.tree_util.tree_map", "jax._src.tree_util.tree_map_with_path", "jax._src.tree_util.tree_reduce", "jax._src.tree_util.tree_reduce_associative", "jax._src.tree_util.tree_structure", "jax._src.tree_util.tree_transpose", "jax._src.tree_util.tree_unflatten", "jax._src.tree_util.treedef_children", "jax._src.tree_util.treedef_is_leaf", "jax._src.tree_util.treedef_tuple" ], "n_typable": 95, "n_typed": 65, "n_any": 24, "n_untyped": 6, "n_functions": 24, "n_function_overloads": 24, "n_function_params": 62, "n_methods": 1, "n_method_overloads": 1, "n_method_params": 3, "n_classes": 1, "n_attrs": 7, "n_properties": 0, "n_type_ignores": 3 }, { "path": "jax/tree_util.py", "symbol_reports": [ { "kind": "attr", "name": "jax.tree_util.PyTreeDef", "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_typable": 1 } ], "type_ignores": [], "name": "jax.tree_util", "names": [ "jax.tree_util.PyTreeDef" ], "n_typable": 1, "n_typed": 0, "n_any": 0, "n_untyped": 1, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 }, { "path": "jax/_src/typing.py", "symbol_reports": [ { "kind": "attr", "name": "jax._src.typing.DTypeLike", "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_typable": 1 } ], "type_ignores": [], "name": "jax._src.typing", "names": [ "jax._src.typing.DTypeLike" ], "n_typable": 1, "n_typed": 1, "n_any": 0, "n_untyped": 0, "n_functions": 0, "n_function_overloads": 0, "n_function_params": 0, "n_methods": 0, "n_method_overloads": 0, "n_method_params": 0, "n_classes": 0, "n_attrs": 1, "n_properties": 0, "n_type_ignores": 0 } ], "typecheckers": {}, "n_modules": 280, "n_typable": 12091, "n_typed": 7898, "n_any": 278, "n_untyped": 3915, "n_functions": 1867, "n_function_overloads": 1945, "n_function_params": 6038, "n_methods": 1264, "n_method_params": 1839, "n_method_overloads": 1264, "n_classes": 343, "n_attrs": 1683, "n_properties": 262, "type_ignores": [ { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": [ "override" ] }, { "kind": "type", "rules": [ "override" ] }, { "kind": "type", "rules": [ "misc" ] }, { "kind": "type", "rules": [ "misc" ] }, { "kind": "type", "rules": [ "misc" ] }, { "kind": "type", "rules": [ "misc" ] }, { "kind": "type", "rules": [ "misc" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": [ "missing-attribute" ] }, { "kind": "type", "rules": [ "missing-attribute" ] }, { "kind": "type", "rules": [ "missing-attribute" ] }, { "kind": "type", "rules": [ "missing-attribute" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "pyrefly", "rules": [ "bad-param-name-override" ] }, { "kind": "pyrefly", "rules": [ "bad-param-name-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-param-name-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": [ "import-not-found" ] }, { "kind": "type", "rules": null }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "type", "rules": null }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": [ "bad-return-type", "unused-ignore" ] }, { "kind": "type", "rules": [ "bad-return-type", "unused-ignore" ] }, { "kind": "type", "rules": null }, { "kind": "pyrefly", "rules": [ "bad-class-definition" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "method-assign" ] }, { "kind": "type", "rules": [ "method-assign" ] }, { "kind": "type", "rules": null }, { "kind": "pyrefly", "rules": [ "bad-assignment" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": null }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "type", "rules": [ "method-assign" ] }, { "kind": "type", "rules": [ "method-assign" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": [ "arg-type" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": [ "missing-attribute" ] }, { "kind": "type", "rules": [ "missing-attribute" ] }, { "kind": "type", "rules": [ "bad-assignment" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": [ "method-assign" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "assignment" ] }, { "kind": "type", "rules": [ "override" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "type", "rules": null }, { "kind": "pyrefly", "rules": [ "bad-override" ] }, { "kind": "pyrefly", "rules": [ "bad-param-name-override" ] }, { "kind": "pyrefly", "rules": [ "bad-param-name-override" ] }, { "kind": "type", "rules": null }, { "kind": "type", "rules": [ "bad-override", "assignment" ] }, { "kind": "pyrefly", "rules": [ "missing-attribute" ] }, { "kind": "pyrefly", "rules": [ "missing-attribute" ] }, { "kind": "pyrefly", "rules": [ "missing-attribute" ] } ], "n_type_ignores": 181 }