message: Special case for one sample categorical so that TPU memory padding doesn't increase the number of samples dim.
diff: @@ -59,14 +59,31 @@ class CategoricalOp : public XlaOpKernel {
  xla::XlaBuilder* builder = ctx->builder();
  + xla::Shape uniform_shape;
  + int class_dimension;
  + if (num_samples > 1) {
  std::array<int64, 3> uniform_shape_array = {
  {batch_size, num_samples, num_classes}};
  xla::PrimitiveType uniform_xla_type;
  OP_REQUIRES_OK(ctx,
  DataTypeToPrimitiveType(input_type(0), &uniform_xla_type));
  - xla::Shape uniform_shape =
  + uniform_shape =
  xla::ShapeUtil::MakeShape(uniform_xla_type, uniform_shape_array);
  - auto uniforms =
  + class_dimension = 2;
  + } else {
  + // Have a special case for when we only need one sample, because
  + // dimensions may be padded on architectures with tiled memory layouts, so
  + // if the num_classes or batch size is large then this can lead to
  + // expensive wasted memory.
  + std::array<int64, 2> uniform_shape_array = {{batch_size, num_classes}};
  + xla::PrimitiveType uniform_xla_type;
  + OP_REQUIRES_OK(ctx,
  + DataTypeToPrimitiveType(input_type(0), &uniform_xla_type));
  + uniform_shape =
  + xla::ShapeUtil::MakeShape(uniform_xla_type, uniform_shape_array);
  + class_dimension = 1;
  + }
  + xla::XlaOp uniforms =
  xla::RngUniform(XlaHelpers::Zero(builder, input_type(0)),
  XlaHelpers::One(builder, input_type(0)), uniform_shape);
  @@ -74,14 +91,18 @@ class CategoricalOp : public XlaOpKernel {
  // See:
  // https://hips.seas.harvard.edu/blog/2013/04/06/the-gumbel-max-trick-for-discrete-distributions/
  // TODO(b/68769470): Switch to using a cumulative sum approach.
  - auto softmax_entries = xla::Sub(logits, xla::Log(-xla::Log(uniforms)),
  - /*broadcast_dimensions=*/{0, 2});
  + auto softmax_entries =
  + xla::Sub(logits, xla::Log(-xla::Log(uniforms)),
  + /*broadcast_dimensions=*/{0, class_dimension});
  xla::PrimitiveType xla_output_type;
  OP_REQUIRES_OK(ctx,
  DataTypeToPrimitiveType(output_type(0), &xla_output_type));
  - xla::XlaOp argmax =
  - XlaHelpers::ArgMax(softmax_entries, xla_output_type, /*axis=*/2);
  + xla::XlaOp argmax = XlaHelpers::ArgMax(softmax_entries, xla_output_type,
  + /*axis=*/class_dimension);
  + if (num_samples == 1) {
  + argmax = xla::Reshape(argmax, {batch_size, 1});
  + }
  ctx->SetOutput(0, argmax);
  }
  