Skip to content

Refresh non-variable reshape expressions in Grappler #21

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
stevenvar wants to merge 5 commits into
base: for-serving-2.20
Choose a base branch
from
Draft

Refresh non-variable reshape expressions in Grappler #21

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
c057b93
Add dynamic size expression tests and fix propagation gaps (#27)
stevenvar Jul 29, 2026
5e677f7
Support max and conditional dynamic shape expressions (#4)
stevenvar Jul 29, 2026
7f3d003
Preserve dynamic shape expressions across XLA shape reconstruction (#13)
stevenvar Jul 29, 2026
61d4725
Reject dynamic shape expressions in MLIR XLA kernels (#19)
stevenvar Jul 29, 2026
01fc0ba
Restore selective reshape expr refresh
stevenvar Apr 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,17 @@ std::string ExprProtoToString(const ExpressionProto& e) {
case ExpressionProto::kDivNode:
return absl::StrCat("(", ExprProtoToString(e.div_node().lhs()), " / ",
ExprProtoToString(e.div_node().rhs()), ")");
case ExpressionProto::kMaxNode:
return absl::StrCat("max(", ExprProtoToString(e.max_node().lhs()), ", ",
ExprProtoToString(e.max_node().rhs()), ")");
case ExpressionProto::kGtNode:
return absl::StrCat("(", ExprProtoToString(e.gt_node().lhs()), " > ",
ExprProtoToString(e.gt_node().rhs()), ")");
case ExpressionProto::kSelectNode:
return absl::StrCat("select(", ExprProtoToString(e.select_node().pred()),
", ", ExprProtoToString(e.select_node().on_true()),
", ", ExprProtoToString(e.select_node().on_false()),
")");
default:
return "<none>";
}
Expand Down
33 changes: 33 additions & 0 deletions tensorflow/compiler/jit/kernels/xla_ops.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -414,6 +414,23 @@ std::unique_ptr<DimExpr> ExprFromProto(const ExpressionProto& proto) {
auto rhs = ExprFromProto(proto.div_node().rhs());
return std::make_unique<ExprDiv>(lhs.release(), rhs.release());
}
case ExpressionProto::kMaxNode: {
auto lhs = ExprFromProto(proto.max_node().lhs());
auto rhs = ExprFromProto(proto.max_node().rhs());
return std::make_unique<ExprMax>(lhs.release(), rhs.release());
}
case ExpressionProto::kGtNode: {
auto lhs = ExprFromProto(proto.gt_node().lhs());
auto rhs = ExprFromProto(proto.gt_node().rhs());
return std::make_unique<ExprGt>(lhs.release(), rhs.release());
}
case ExpressionProto::kSelectNode: {
auto pred = ExprFromProto(proto.select_node().pred());
auto on_true = ExprFromProto(proto.select_node().on_true());
auto on_false = ExprFromProto(proto.select_node().on_false());
return std::make_unique<ExprSelect>(pred.release(), on_true.release(),
on_false.release());
}
case ExpressionProto::NODE_TYPE_NOT_SET:
default:
return nullptr;
Expand Down Expand Up @@ -445,6 +462,22 @@ static xla::DExpr DimExprToDExpr(const DimExpr* e) {
auto* ee = static_cast<const ExprDiv*>(e);
return DimExprToDExpr(ee->lhs()) / DimExprToDExpr(ee->rhs());
}
case DimExpr::Kind::kMax: {
auto* ee = static_cast<const ExprMax*>(e);
return xla::DExpr::Max(DimExprToDExpr(ee->lhs()),
DimExprToDExpr(ee->rhs()));
}
case DimExpr::Kind::kGt: {
auto* ee = static_cast<const ExprGt*>(e);
return xla::DExpr::Gt(DimExprToDExpr(ee->lhs()),
DimExprToDExpr(ee->rhs()));
}
case DimExpr::Kind::kSelect: {
auto* ee = static_cast<const ExprSelect*>(e);
return xla::DExpr::Select(DimExprToDExpr(ee->pred()),
DimExprToDExpr(ee->on_true()),
DimExprToDExpr(ee->on_false()));
}
}
return xla::DExpr::Unknown();
}
Expand Down
44 changes: 44 additions & 0 deletions tensorflow/compiler/jit/mark_for_compilation_pass.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -714,6 +714,17 @@ std::string ExprProtoToString(const ExpressionProto& e) {
case ExpressionProto::kDivNode:
return absl::StrCat("(", ExprProtoToString(e.div_node().lhs()), " / ",
ExprProtoToString(e.div_node().rhs()), ")");
case ExpressionProto::kMaxNode:
return absl::StrCat("max(", ExprProtoToString(e.max_node().lhs()), ", ",
ExprProtoToString(e.max_node().rhs()), ")");
case ExpressionProto::kGtNode:
return absl::StrCat("(", ExprProtoToString(e.gt_node().lhs()), " > ",
ExprProtoToString(e.gt_node().rhs()), ")");
case ExpressionProto::kSelectNode:
return absl::StrCat("select(", ExprProtoToString(e.select_node().pred()),
", ", ExprProtoToString(e.select_node().on_true()),
", ", ExprProtoToString(e.select_node().on_false()),
")");
default:
return "<none>";
}
Expand Down Expand Up @@ -747,6 +758,23 @@ std::unique_ptr<DimExpr> ExprFromProto(const ExpressionProto& proto) {
auto rhs = ExprFromProto(proto.div_node().rhs());
return std::make_unique<ExprDiv>(lhs.release(), rhs.release());
}
case ExpressionProto::kMaxNode: {
auto lhs = ExprFromProto(proto.max_node().lhs());
auto rhs = ExprFromProto(proto.max_node().rhs());
return std::make_unique<ExprMax>(lhs.release(), rhs.release());
}
case ExpressionProto::kGtNode: {
auto lhs = ExprFromProto(proto.gt_node().lhs());
auto rhs = ExprFromProto(proto.gt_node().rhs());
return std::make_unique<ExprGt>(lhs.release(), rhs.release());
}
case ExpressionProto::kSelectNode: {
auto pred = ExprFromProto(proto.select_node().pred());
auto on_true = ExprFromProto(proto.select_node().on_true());
auto on_false = ExprFromProto(proto.select_node().on_false());
return std::make_unique<ExprSelect>(pred.release(), on_true.release(),
on_false.release());
}
case ExpressionProto::NODE_TYPE_NOT_SET:
default:
return nullptr;
Expand Down Expand Up @@ -779,6 +807,22 @@ static xla::DExpr DimExprToDExpr(const DimExpr* e) {
auto* ee = static_cast<const ExprDiv*>(e);
return DimExprToDExpr(ee->lhs()) / DimExprToDExpr(ee->rhs());
}
case DimExpr::Kind::kMax: {
auto* ee = static_cast<const ExprMax*>(e);
return xla::DExpr::Max(DimExprToDExpr(ee->lhs()),
DimExprToDExpr(ee->rhs()));
}
case DimExpr::Kind::kGt: {
auto* ee = static_cast<const ExprGt*>(e);
return xla::DExpr::Gt(DimExprToDExpr(ee->lhs()),
DimExprToDExpr(ee->rhs()));
}
case DimExpr::Kind::kSelect: {
auto* ee = static_cast<const ExprSelect*>(e);
return xla::DExpr::Select(DimExprToDExpr(ee->pred()),
DimExprToDExpr(ee->on_true()),
DimExprToDExpr(ee->on_false()));
}
}
return xla::DExpr();
}
Expand Down
32 changes: 28 additions & 4 deletions tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,8 @@ void BatchToSpace(XlaOpKernelContext* ctx, const xla::XlaOp input,
const int input_rank = input_tensor_shape.dims();
const absl::InlinedVector<int64_t, 4> input_shape =
input_tensor_shape.dim_sizes();
const std::vector<xla::DExpr> input_exprs =
input_tensor_shape.get_filled_expressions();
const int block_rank = block_shape.size();

OP_REQUIRES(
Expand Down Expand Up @@ -76,11 +78,21 @@ void BatchToSpace(XlaOpKernelContext* ctx, const xla::XlaOp input,
") is not divisible by product of block sizes (",
block_num_elems, ")"));
std::vector<int64_t> reshaped_shape(input_rank + block_rank);
std::vector<xla::DExpr> reshaped_exprs(input_rank + block_rank);
std::copy(block_shape.begin(), block_shape.end(), reshaped_shape.begin());
std::fill(reshaped_exprs.begin(), reshaped_exprs.begin() + block_rank,
xla::DExpr::Const(0));
for (int i = 0; i < block_rank; ++i) {
reshaped_exprs[i] = xla::DExpr::Const(block_shape[i]);
}
reshaped_shape[block_rank] = batch_size / block_num_elems;
reshaped_exprs[block_rank] =
(input_exprs[0] / xla::DExpr::Const(block_num_elems)).simplify();
std::copy(input_shape.begin() + 1, input_shape.end(),
reshaped_shape.begin() + block_rank + 1);
xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape);
std::copy(input_exprs.begin() + 1, input_exprs.end(),
reshaped_exprs.begin() + block_rank + 1);
xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape, reshaped_exprs);

// 2. Permute dimensions of `reshaped` to produce `permuted` of shape
// [batch / prod(block_shape),
Expand Down Expand Up @@ -111,15 +123,22 @@ void BatchToSpace(XlaOpKernelContext* ctx, const xla::XlaOp input,
// ...,
// input_shape[N-1]]
std::vector<int64_t> reshaped_permuted_shape(input_rank);
std::vector<xla::DExpr> reshaped_permuted_exprs(input_rank);
reshaped_permuted_shape[0] = batch_size / block_num_elems;
reshaped_permuted_exprs[0] =
(input_exprs[0] / xla::DExpr::Const(block_num_elems)).simplify();
for (int i = 0; i < block_rank; ++i) {
reshaped_permuted_shape[1 + i] = block_shape[i] * input_shape[1 + i];
reshaped_permuted_exprs[1 + i] =
(xla::DExpr::Const(block_shape[i]) * input_exprs[1 + i]).simplify();
}
std::copy(remainder_shape.begin(), remainder_shape.end(),
reshaped_permuted_shape.begin() + 1 + block_rank);
std::copy(input_exprs.begin() + 1 + block_rank, input_exprs.end(),
reshaped_permuted_exprs.begin() + 1 + block_rank);

xla::XlaOp reshaped_permuted =
xla::Reshape(permuted, reshaped_permuted_shape);
xla::Reshape(permuted, reshaped_permuted_shape, reshaped_permuted_exprs);

// 4. Crop the start and end of dimensions `[1, ..., M]` of
// `reshaped_permuted` according to `crops` to produce the output of shape:
Expand All @@ -133,21 +152,26 @@ void BatchToSpace(XlaOpKernelContext* ctx, const xla::XlaOp input,
std::vector<int64_t> start_indices(input_rank, 0);
std::vector<int64_t> end_indices = reshaped_permuted_shape;
std::vector<int64_t> strides(input_rank, 1);
std::vector<xla::DExpr> start_exprs(input_rank, xla::DExpr::Const(0));
std::vector<xla::DExpr> end_exprs(reshaped_permuted_exprs.begin(),
reshaped_permuted_exprs.end());
for (int i = 0; i < block_rank; ++i) {
int64_t crop_start = crops.Get<int64_t>({i, 0});
int64_t crop_end = crops.Get<int64_t>({i, 1});
OP_REQUIRES(ctx, crop_start >= 0 && crop_end >= 0,
errors::InvalidArgument("Crops must be non-negative"));
start_indices[1 + i] = crop_start;
end_indices[1 + i] -= crop_end;
start_exprs[1 + i] = xla::DExpr::Const(crop_start);
end_exprs[1 + i] = (reshaped_permuted_exprs[1 + i] - crop_end).simplify();
OP_REQUIRES(
ctx, start_indices[1 + i] <= end_indices[1 + i],
errors::InvalidArgument(
"Cropped size must be non-negative: start: ", crop_start,
" end: ", crop_end, " size ", reshaped_permuted_shape[1 + i]));
}
xla::XlaOp output =
xla::Slice(reshaped_permuted, start_indices, end_indices, strides);
xla::XlaOp output = xla::Slice(reshaped_permuted, start_indices, end_indices,
start_exprs, end_exprs, strides);
ctx->SetOutput(0, output);
}

Expand Down
13 changes: 9 additions & 4 deletions tensorflow/compiler/tf2xla/kernels/bincount_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -110,11 +110,15 @@ class DenseBincountOp : public XlaOpKernel {
scatter_dnums.add_scatter_dims_to_operand_dims(0);

if (rank == 2) {
output_shape = xla::ShapeUtil::MakeShape(dtype, {size, output_size});
output_shape = xla::ShapeUtil::MakeShape(
dtype, {size, output_size},
std::vector<xla::DExpr>{input_shape.expressions(0),
xla::DExpr::Const(output_size)});
scatter_dnums.add_inserted_window_dims(1);
scatter_dnums.add_scatter_dims_to_operand_dims(1);
auto i_shape =
xla::ShapeUtil::MakeShape(input_xla_type, {input_shape.dimensions()});
auto i_shape = xla::ShapeUtil::MakeShape(input_xla_type,
input_shape.dimensions(),
input_shape.expressions());
auto i = xla::Iota(ctx->builder(), i_shape, 0);
xla::DExpr flattened_expr =
input_shape.expressions(0) * input_shape.expressions(1);
Expand All @@ -131,7 +135,8 @@ class DenseBincountOp : public XlaOpKernel {
updates = xla::Broadcast(
one, {input_shape.dimensions(0) * input_shape.dimensions(1)});
output = xla::Broadcast(
zero, {output_shape.dimensions(0), output_shape.dimensions(1)});
zero, {output_shape.dimensions(0), output_shape.dimensions(1)},
{output_shape.expressions(0), output_shape.expressions(1)});
if (has_weights && !binary_output_) {
weights = xla::Reshape(
weights, {input_shape.dimensions(0) * input_shape.dimensions(1)},
Expand Down
26 changes: 26 additions & 0 deletions tensorflow/compiler/tf2xla/kernels/const_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,16 @@ bool IsDynamicExpressionProto(const ExpressionProto& proto) {
case ExpressionProto::kDivNode:
return IsDynamicExpressionProto(proto.div_node().lhs()) ||
IsDynamicExpressionProto(proto.div_node().rhs());
case ExpressionProto::kMaxNode:
return IsDynamicExpressionProto(proto.max_node().lhs()) ||
IsDynamicExpressionProto(proto.max_node().rhs());
case ExpressionProto::kGtNode:
return IsDynamicExpressionProto(proto.gt_node().lhs()) ||
IsDynamicExpressionProto(proto.gt_node().rhs());
case ExpressionProto::kSelectNode:
return IsDynamicExpressionProto(proto.select_node().pred()) ||
IsDynamicExpressionProto(proto.select_node().on_true()) ||
IsDynamicExpressionProto(proto.select_node().on_false());
case ExpressionProto::kConstantValue:
case ExpressionProto::NODE_TYPE_NOT_SET:
return false;
Expand Down Expand Up @@ -158,6 +168,22 @@ static xla::DExpr DimExprToDExpr(const DimExpr* e) {
const auto* ee = static_cast<const ExprDiv*>(e);
return DimExprToDExpr(ee->lhs()) / DimExprToDExpr(ee->rhs());
}
case DimExpr::Kind::kMax: {
const auto* ee = static_cast<const ExprMax*>(e);
return xla::DExpr::Max(DimExprToDExpr(ee->lhs()),
DimExprToDExpr(ee->rhs()));
}
case DimExpr::Kind::kGt: {
auto* ee = static_cast<const ExprGt*>(e);
return xla::DExpr::Gt(DimExprToDExpr(ee->lhs()),
DimExprToDExpr(ee->rhs()));
}
case DimExpr::Kind::kSelect: {
auto* ee = static_cast<const ExprSelect*>(e);
return xla::DExpr::Select(DimExprToDExpr(ee->pred()),
DimExprToDExpr(ee->on_true()),
DimExprToDExpr(ee->on_false()));
}
}
return xla::DExpr();
}
Expand Down
41 changes: 39 additions & 2 deletions tensorflow/compiler/tf2xla/kernels/depthtospace_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ class DepthToSpaceOp : public XlaOpKernel {
OP_REQUIRES_OK(ctx, input_xla_shape.status());
absl::Span<const int64_t> input_shape =
input_xla_shape.value().dimensions();
const xla::Shape& input_shape_with_exprs = input_xla_shape.value();
int input_rank = input_shape.size();

static const int kRequiredDims = 4;
Expand All @@ -77,20 +78,31 @@ class DepthToSpaceOp : public XlaOpKernel {
std::vector<int64_t> reshaped_shape;
std::vector<int64_t> transpose_order;
std::vector<int64_t> output_shape;
std::vector<xla::DExpr> reshaped_exprs;
std::vector<xla::DExpr> output_exprs;
reshaped_shape.reserve(input_rank);
transpose_order.reserve(input_rank);
output_shape.reserve(input_rank);
reshaped_exprs.reserve(input_rank + num_spatial_dims);
output_exprs.reserve(input_rank);
if (data_format == FORMAT_NHWC) {
reshaped_shape.push_back(input_shape[0]);
reshaped_exprs.push_back(input_shape_with_exprs.expressions(0));
for (int i = 0; i < num_spatial_dims; ++i) {
reshaped_shape.push_back(input_shape[1 + i]);
reshaped_exprs.push_back(input_shape_with_exprs.expressions(1 + i));
}
int64_t block_elems = 1;
for (int i = 0; i < num_spatial_dims; ++i) {
reshaped_shape.push_back(block_size_);
reshaped_exprs.push_back(xla::DExpr::Const(block_size_));
block_elems *= block_size_;
}
reshaped_shape.push_back(input_shape[feature_dim] / block_elems);
reshaped_exprs.push_back(
(input_shape_with_exprs.expressions(feature_dim) /
xla::DExpr::Const(block_elems))
.simplify());

transpose_order.push_back(0);
for (int i = 0; i < num_spatial_dims; ++i) {
Expand All @@ -100,21 +112,37 @@ class DepthToSpaceOp : public XlaOpKernel {
transpose_order.push_back(feature_dim + num_spatial_dims);

output_shape.push_back(input_shape[0]);
output_exprs.push_back(input_shape_with_exprs.expressions(0));
for (int i = 0; i < num_spatial_dims; ++i) {
output_shape.push_back(input_shape[1 + i] * block_size_);
output_exprs.push_back(
(input_shape_with_exprs.expressions(1 + i) *
xla::DExpr::Const(block_size_))
.simplify());
}
output_shape.push_back(input_shape[feature_dim] / block_elems);
output_exprs.push_back(
(input_shape_with_exprs.expressions(feature_dim) /
xla::DExpr::Const(block_elems))
.simplify());
} else {
// NCHW format.
reshaped_shape.push_back(input_shape[0]);
reshaped_exprs.push_back(input_shape_with_exprs.expressions(0));
int64_t block_elems = 1;
for (int i = 0; i < num_spatial_dims; ++i) {
reshaped_shape.push_back(block_size_);
reshaped_exprs.push_back(xla::DExpr::Const(block_size_));
block_elems *= block_size_;
}
reshaped_shape.push_back(input_shape[feature_dim] / block_elems);
reshaped_exprs.push_back(
(input_shape_with_exprs.expressions(feature_dim) /
xla::DExpr::Const(block_elems))
.simplify());
for (int i = 0; i < num_spatial_dims; ++i) {
reshaped_shape.push_back(input_shape[2 + i]);
reshaped_exprs.push_back(input_shape_with_exprs.expressions(2 + i));
}

transpose_order.push_back(0);
Expand All @@ -125,9 +153,18 @@ class DepthToSpaceOp : public XlaOpKernel {
}

output_shape.push_back(input_shape[0]);
output_exprs.push_back(input_shape_with_exprs.expressions(0));
output_shape.push_back(input_shape[feature_dim] / block_elems);
output_exprs.push_back(
(input_shape_with_exprs.expressions(feature_dim) /
xla::DExpr::Const(block_elems))
.simplify());
for (int i = 0; i < num_spatial_dims; ++i) {
output_shape.push_back(input_shape[2 + i] * block_size_);
output_exprs.push_back(
(input_shape_with_exprs.expressions(2 + i) *
xla::DExpr::Const(block_size_))
.simplify());
}
}

Expand All @@ -148,7 +185,7 @@ class DepthToSpaceOp : public XlaOpKernel {
") is not divisible by square of the block size (",
block_size_, ")"));

xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape);
xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape, reshaped_exprs);

// 2. Permute dimensions of `reshaped` to produce
// `permuted_reshaped` of shape:
Expand All @@ -169,7 +206,7 @@ class DepthToSpaceOp : public XlaOpKernel {
// input_shape[2] * block_size_,
// depth / (block_size_ * block_size_)]
//
xla::XlaOp output = xla::Reshape(permuted_reshaped, output_shape);
xla::XlaOp output = xla::Reshape(permuted_reshaped, output_shape, output_exprs);

// If this used to be a vectorized format turn it back now.
if (data_format != data_format_) {
Expand Down
Loading