Commit 368674df authored by Smit Hinsu's avatar Smit Hinsu Committed by TensorFlower Gardener
Browse files

PR #25565: TFTRT: Support GatherV2 op

Please approve this CL. It will be submitted automatically, and its GitHub pull request will be marked as merged.

Imported from GitHub PR #25565

* Add support for GatherV2 Op.
* Add unit tests
* Add CheckAxis helper method (need to use in other functions, will issue separate PR for this)

https://www.tensorflow.org/api_docs/cc/class/tensorflow/ops/gather-v2

The "scalar indices" case is currently unsupported because we don't allow scalar tensors, but they are now supported in TRT 5.0+ or maybe earlier, need to double check (TRT calls these 0-D tensors). We will have to update our code to allow these tensors in order to allow this case to convert.

Will add support for constant (weight) indices with automatic weight -> tensor helper soon.

Copybara import of the project:

  - bad24cee6eb7888912d8a582ae75981ecc893dce Support GatherV2 op. Add unit tests. Add CheckAxis helper... by Trevor Morris <tmorris@nvidia.com>
  - bb1597436974e6c840d2da3af30292877f51e062 Fix formatting and add TODO for datatypes by Trevor Morris <tmorris@nvidia.com>
  - 1899c671d97b7efcf1adb87a3f32158fd4603fcd Make templated test function for Gather to test all appli... by Trevor Morris <tmorris@nvidia.com>
  - 48e7536c9a81415ac7487c68807e25d2c92a7a2c Improve ConvertAxis. Change TRT_ShapedWeights::ToVector -... by Trevor Morris <tmorris@nvidia.com>
  - d5c684610c42b8c7e02112d270fd12afed2ed484 Remove unnecessary std::to_string by Trevor Morris <tmorris@nvidia.com>
  - 6a34e557898ef1b71b3c78c18a929ffb7d11c2d4 Use absl::Span<const T> by Trevor Morris <tmorris@nvidia.com>
  - ff7a35e6e131cdca7082af1b274b406545ece6fa Remove to_string. Remove unused check_bounds arg by Trevor Morris <tmorris@nvidia.com>
  - 90d8c78e62b39eedadb9f3c8a6367fda6cf9120e Merge ff7a35e6e131cdca7082af1b274b406545ece6fa into 39a64... by Trevor Morris <trevoraidanmorris@gmail.com>

COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/tensorflow/pull/25565 from trevor-m:tmorris_tftrt_gather ff7a35e6e131cdca7082af1b274b406545ece6fa
PiperOrigin-RevId: 233892061
parent 71a9efdc
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment