Searched refs:ArgSignature (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | jax_jit.h | 31 struct ArgSignature { struct
32 ArgSignature(xla::PrimitiveType dtype, absl::Span<const xla::int64> shape, in ArgSignature() function
41 bool operator==(const ArgSignature& other) const { argument
45 bool operator!=(const ArgSignature& other) const { return !(*this == other); }
50 H AbslHashValue(H h, const ArgSignature& s) { in AbslHashValue()
87 std::vector<ArgSignature> dynamic_args_signatures;
156 xla::StatusOr<ArgSignature> ArgSignatureOfValue(pybind11::handle arg,
|
| D | jax_jit.cc | 65 std::string ArgSignature::DebugString() const { in DebugString()
378 std::function<xla::StatusOr<ArgSignature>(py::handle, bool)>;
381 xla::StatusOr<ArgSignature> ArgSignatureOfValue(pybind11::handle arg, in ArgSignatureOfValue()
394 [](py::handle, bool) -> xla::StatusOr<ArgSignature> { in ArgSignatureOfValue()
395 return ArgSignature(xla::PrimitiveType::PRED, {}, true); in ArgSignatureOfValue()
398 [](py::handle h, bool jax_enable_x64) -> xla::StatusOr<ArgSignature> { in ArgSignatureOfValue()
400 return ArgSignature(xla::PrimitiveType::S64, {}, true); in ArgSignatureOfValue()
402 return ArgSignature(xla::PrimitiveType::S32, {}, true); in ArgSignatureOfValue()
407 bool jax_enable_x64) -> xla::StatusOr<ArgSignature> { in ArgSignatureOfValue()
411 return ArgSignature(xla::PrimitiveType::F64, {}, weak_type); in ArgSignatureOfValue()
[all …]
|