Searched refs:PyArgSignature (Results 1 – 4 of 4) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | py_values.cc | 351 std::string PyArgSignature::DebugString() const { in DebugString()
362 std::function<StatusOr<PyArgSignature>(py::handle, bool)>;
364 StatusOr<PyArgSignature> PyArgSignatureOfValue(py::handle arg, in PyArgSignatureOfValue()
374 [](py::handle, bool) -> StatusOr<PyArgSignature> { in PyArgSignatureOfValue()
375 return PyArgSignature(PrimitiveType::PRED, {}, true); in PyArgSignatureOfValue()
378 [](py::handle h, bool jax_enable_x64) -> StatusOr<PyArgSignature> { in PyArgSignatureOfValue()
381 return PyArgSignature(PrimitiveType::S64, {}, true); in PyArgSignatureOfValue()
383 return PyArgSignature(PrimitiveType::S32, {}, true); in PyArgSignatureOfValue()
388 bool jax_enable_x64) -> StatusOr<PyArgSignature> { in PyArgSignatureOfValue()
392 return PyArgSignature(PrimitiveType::F64, {}, weak_type); in PyArgSignatureOfValue()
[all …]
|
| D | py_values.h | 71 struct PyArgSignature { struct
72 PyArgSignature(PrimitiveType dtype, absl::Span<const int64_t> shape, in PyArgSignature() argument
81 bool operator==(const PyArgSignature& other) const { argument
85 bool operator!=(const PyArgSignature& other) const {
93 StatusOr<PyArgSignature> PyArgSignatureOfValue(pybind11::handle arg,
97 H AbslHashValue(H h, const xla::PyArgSignature& s) { in AbslHashValue()
|
| D | jax_jit.h | 103 absl::InlinedVector<xla::PyArgSignature, 2> dynamic_arg_signatures;
|
| D | jax_jit.cc | 137 const xla::PyArgSignature& s) { in DebugString()
1451 py::class_<xla::PyArgSignature> arg_signature(jitlib, "PyArgSignature"); in BuildJaxjitSubmodule()
1454 [](const xla::PyArgSignature& sig) { in BuildJaxjitSubmodule()
1459 [](const xla::PyArgSignature& sig) { in BuildJaxjitSubmodule()
1462 .def_readonly("weak_type", &xla::PyArgSignature::weak_type); in BuildJaxjitSubmodule()
|