Searched refs:jax_enable_x64_ (Results 1 – 2 of 2) sorted by relevance
186 absl::optional<bool> jax_enable_x64_ = absl::nullopt; member in jax::PmapFunction260 if (jax_enable_x64_ == absl::nullopt) { in Call()261 jax_enable_x64_ = py::cast<bool>(get_jax_enable_x64_()); in Call()271 auto signature_or_error = ArgSignatureOfValue(arg, jax_enable_x64_.value()); in Call()
900 absl::optional<bool> jax_enable_x64_ = absl::nullopt; member in jax::__anone04173081411::CompiledFunction1118 jax_enable_x64_ = py::cast<bool>(get_jax_enable_x64_()); in Call()1152 if (!ConvertArgsToBuffers(jax_enable_x64_.value(), *default_pyclient_, in Call()