Method Compile

tensorflow/compiler/tf2xla/kernels/transpose_op.cc:39–93

      : XlaOpKernel(ctx), conjugate_(conjugate) {}

  void Compile(XlaOpKernelContext* ctx) override {
    const TensorShape input_shape = ctx->InputShape("x");
    const TensorShape perm_tensor_shape = ctx->InputShape("perm");

    // Preliminary validation of sizes.
    OP_REQUIRES(ctx, TensorShapeUtils::IsVector(perm_tensor_shape),
                errors::InvalidArgument("perm must be a vector, not ",
                                        perm_tensor_shape.DebugString()));

    const int dims = input_shape.dims();
    OP_REQUIRES(ctx, dims == perm_tensor_shape.num_elements(),
                errors::InvalidArgument("transpose expects a vector of size ",
                                        input_shape.dims(),
                                        ". But input(1) is a vector of size ",
                                        perm_tensor_shape.num_elements()));

    std::vector<int64> perm;
    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector("perm", &perm));

    std::vector<int64> transposed_order;
    // Check whether permutation is a permutation of integers of [0 .. dims).
    absl::InlinedVector<bool, 8> bits(dims);
    bool is_identity = true;
    for (int i = 0; i < dims; ++i) {
      const int64 d = perm[i];
      OP_REQUIRES(
          ctx, 0 <= d && d < dims,
          errors::InvalidArgument(d, " is out of range [0 .. ", dims, ")"));
      bits[d] = true;
      transposed_order.push_back(d);
      if (d != i) {
        is_identity = false;
      }
    }
    for (int i = 0; i < dims; ++i) {
      OP_REQUIRES(
          ctx, bits[i],
          errors::InvalidArgument(i, " is missing from 'perm' argument."));
    }

    xla::XlaOp transposed;
    // 0-D, 1-D, and identity transposes do nothing.
    if (dims <= 1 || is_identity) {
      transposed = ctx->Input("x");
    } else {
      transposed = xla::Transpose(ctx->Input("x"), transposed_order);
    }

    // Conjugate the transposed result if this is ConjugateTransposeOp.
    if (conjugate_) {
      ctx->SetOutput(0, xla::Conj(transposed));
    } else {
      ctx->SetOutput(0, transposed);
    }
  }

 private:
  const bool conjugate_;