Searched defs:static_argnums (Results 1 – 2 of 2) sorted by relevance
196 std::vector<int> static_argnums, in PmapFunction()231 const std::vector<int>& static_argnums() const { return static_argnums_; } in static_argnums() function in jax::PmapFunction710 std::vector<int> static_argnums, in InitializePmapFunction()720 std::vector<int> static_argnums, in MakePmapFunction()901 std::vector<int> static_argnums = in BuildPmapSubmodule() local952 py::function python_shard_arg_fallback) -> py::object { in BuildPmapSubmodule()
235 absl::Span<int const> static_argnums, in ParseArguments()509 const std::vector<int>& static_argnums() const { return static_argnums_; } in static_argnums() function in jax::__anon838f46370811::CompiledFunction608 std::vector<int> static_argnums, in CompiledFunction()1203 std::vector<int> static_argnums, in InitializeCompiledFunction()1218 std::vector<int> static_argnums, in MakeCompiledFunction()1360 std::vector<int> static_argnums = in BuildJaxjitSubmodule() local1401 std::shared_ptr<CompiledFunctionCache> cache) -> py::object { in BuildJaxjitSubmodule()