diff --git a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.cpp b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.cpp
index 0fc10cf..26dd4e9 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.cpp
@@ -191,6 +191,51 @@
 }
 
 //===----------------------------------------------------------------------===//
+// custom<ShapedOperandList>($values, type($values), $value_dims)
+//===----------------------------------------------------------------------===//
+// %value : type{%dynamic_dims}, ...
+
+static ParseResult parseShapedOperandList(
+    OpAsmParser &parser,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
+    SmallVectorImpl<Type> &valueTypes,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &valueDims) {
+  do {
+    values.emplace_back();
+    valueTypes.emplace_back();
+    if (failed(parser.parseOperand(values.back())) ||
+        failed(parser.parseColon()) ||
+        failed(parser.parseType(valueTypes.back())))
+      return failure();
+    if (int64_t dynamicDimCount =
+            cast<ShapedType>(valueTypes.back()).getNumDynamicDims()) {
+      if (failed(parser.parseOperandList(valueDims, dynamicDimCount,
+                                         AsmParser::Delimiter::Braces)))
+        return failure();
+    }
+  } while (succeeded(parser.parseOptionalComma()));
+  return success();
+}
+
+static void printShapedOperandList(OpAsmPrinter &p, Operation *op,
+                                   ValueRange values, TypeRange valueTypes,
+                                   ValueRange valueDims) {
+  llvm::interleaveComma(llvm::zip_equal(values, valueTypes), p, [&](auto it) {
+    auto [value, valueType] = it;
+    p << value;
+    p << " : ";
+    p << valueType;
+    if (int64_t dynamicDimCount =
+            cast<ShapedType>(valueType).getNumDynamicDims()) {
+      p << "{";
+      llvm::interleaveComma(valueDims.take_front(dynamicDimCount), p);
+      valueDims = valueDims.drop_front(dynamicDimCount);
+      p << "}";
+    }
+  });
+}
+
+//===----------------------------------------------------------------------===//
 // custom<WorkgroupCountRegion>($body)
 //===----------------------------------------------------------------------===//
 
@@ -1649,6 +1694,35 @@
 }
 
 //===----------------------------------------------------------------------===//
+// flow.tensor.trace
+//===----------------------------------------------------------------------===//
+
+LogicalResult TensorTraceOp::verify() {
+  TensorTraceOp op = *this;
+  if (failed(verifyOpDynamicDims(op, op.getValues(), op.getValueDims()))) {
+    return failure();
+  }
+  return success();
+}
+
+ValueRange TensorTraceOp::getOperandDynamicDims(unsigned idx) {
+  auto valueDims = getValueDims();
+  for (unsigned i = 0; i <= idx; ++i) {
+    auto valueType = cast<ShapedType>(getValues()[i].getType());
+    int64_t dynamicDimCount = valueType.getNumDynamicDims();
+    if (i == idx) {
+      return valueDims.take_front(dynamicDimCount);
+    }
+    valueDims = valueDims.drop_front(dynamicDimCount);
+  }
+  return ValueRange{};
+}
+
+ValueRange TensorTraceOp::getResultDynamicDims(unsigned idx) {
+  return ValueRange{};
+}
+
+//===----------------------------------------------------------------------===//
 // Public methods
 //===----------------------------------------------------------------------===//
 
diff --git a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.td b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.td
index c15da61..e3d976c 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.td
+++ b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.td
@@ -1509,21 +1509,37 @@
   let hasFolder = 1;
 }
 
-def FLOW_TensorTraceOp : FLOW_Op<"tensor.trace", []> {
-  let summary = [{trace value(s) operation}];
+def FLOW_TensorTraceOp : FLOW_Op<"tensor.trace", [
+  AttrSizedOperandSegments,
+  DeclareOpInterfaceMethods<Util_ShapeAwareOp>,
+]> {
+  let summary = [{traces one or more tensor values at runtime}];
   let description = [{
     Traces out to a runtime trace sink (console, log file, etc) the given
-    tensors and titles them with the given key. The key is informational only
-    and useful for titling/marking specific sets of tensors for easier
-    searching.
+    tensors. The key is arbitrary and can be used for identifying the set of
+    values being traced.
   }];
 
   let arguments = (ins
     StrAttr:$key,
-    Variadic<FLOW_Tensor>:$operands
+    Variadic<FLOW_Tensor>:$values,
+    FLOW_ShapeDynamicDims:$value_dims
   );
 
-  let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
+  let assemblyFormat = [{
+    $key `=` `[`
+    custom<ShapedOperandList>($values, type($values), $value_dims)
+    `]` attr-dict-with-keyword
+  }];
+
+  let builders = [
+    OpBuilder<(ins "StringRef":$key, "ValueRange":$values), [{
+      build($_builder, $_state, key, values,
+            IREE::Util::buildDynamicDimsForValues($_state.location, values, $_builder));
+    }]>,
+  ];
+
+  let hasVerifier = 1;
 }
 
 //===---------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_ops.mlir b/compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_ops.mlir
index 636b535..6cc909d 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_ops.mlir
+++ b/compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_ops.mlir
@@ -180,3 +180,19 @@
   %0 = flow.tensor.update %arg0, %arg1[%arg2, %arg3] : tensor<?x?xf32>{%c1, %c2} -> %arg1 as tensor<?x4xf32>{%c3}
   return %0 : tensor<?x4xf32>
 }
+
+// -----
+
+// CHECK-LABEL: @tensorTrace
+//  CHECK-SAME: (%[[TENSOR0:.+]]: tensor<5xf32>, %[[TENSOR1:.+]]: tensor<?x3x?xi32>, %[[TENSOR1_DIM0:.+]]: index, %[[TENSOR1_DIM2:.+]]: index)
+func.func @tensorTrace(%tensor0: tensor<5xf32>, %tensor1: tensor<?x3x?xi32>, %tensor1_dim0: index, %tensor1_dim2: index) {
+  //      CHECK: flow.tensor.trace "FOOBAR" = [
+  // CHECK-SAME:   %[[TENSOR0]] : tensor<5xf32>,
+  // CHECK-SAME:   %[[TENSOR1]] : tensor<?x3x?xi32>{%[[TENSOR1_DIM0]], %[[TENSOR1_DIM2]]}
+  // CHECK-SAME: ]
+  flow.tensor.trace "FOOBAR" = [
+    %tensor0 : tensor<5xf32>,
+    %tensor1 : tensor<?x3x?xi32>{%tensor1_dim0, %tensor1_dim2}
+  ]
+  return
+}
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/InjectDispatchTracing.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/InjectDispatchTracing.cpp
index 22c9c79..c70620b 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/InjectDispatchTracing.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/InjectDispatchTracing.cpp
@@ -46,14 +46,14 @@
 
       // Input tensors:
       OpBuilder builder(dispatchOp);
-      builder.create<TensorTraceOp>(
+      builder.create<IREE::Flow::TensorTraceOp>(
           dispatchOp.getLoc(),
           builder.getStringAttr(entryPointName + " inputs"),
           filterTensorValues(dispatchOp.getArguments()));
 
       // Output tensors:
       builder.setInsertionPointAfter(dispatchOp);
-      builder.create<TensorTraceOp>(
+      builder.create<IREE::Flow::TensorTraceOp>(
           dispatchOp.getLoc(),
           builder.getStringAttr(entryPointName + " outputs"),
           filterTensorValues(dispatchOp.getResults()));
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/InsertDispatchDebugTargets.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/InsertDispatchDebugTargets.cpp
index 8dbb63d..6186fd6 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/InsertDispatchDebugTargets.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/InsertDispatchDebugTargets.cpp
@@ -55,19 +55,20 @@
 }
 
 // Inserts flow.tensor.trace ops around the specified dispatch op.
-static void traceOpWithName(DispatchOp dispatchOp, std::string name) {
+static void traceOpWithName(IREE::Flow::DispatchOp dispatchOp,
+                            std::string name) {
   OpBuilder builder(dispatchOp);
   // Input tensors:
-  builder.create<TensorTraceOp>(
+  builder.create<IREE::Flow::TensorTraceOp>(
       dispatchOp.getLoc(), builder.getStringAttr(name + " inputs"),
       filterNonTensorValues(dispatchOp.getArguments()));
 
   // Output tensors:
   OpBuilder::InsertionGuard guard(builder);
   builder.setInsertionPointAfter(dispatchOp);
-  builder.create<TensorTraceOp>(dispatchOp.getLoc(),
-                                builder.getStringAttr(name + " outputs"),
-                                filterNonTensorValues(dispatchOp.getResults()));
+  builder.create<IREE::Flow::TensorTraceOp>(
+      dispatchOp.getLoc(), builder.getStringAttr(name + " outputs"),
+      filterNonTensorValues(dispatchOp.getResults()));
 }
 
 // Breaks the given function on the specified op by simply returning immediately
@@ -162,7 +163,8 @@
         localTraceOrdinal = traceOrdinal;
 
       auto &bodyRegion = op.getFunctionBody();
-      auto dispatchOps = llvm::to_vector<8>(bodyRegion.getOps<DispatchOp>());
+      auto dispatchOps =
+          llvm::to_vector<8>(bodyRegion.getOps<IREE::Flow::DispatchOp>());
 
       // Trace on a valid ordinal.
       if (localTraceOrdinal >= 0 && localTraceOrdinal < dispatchOps.size()) {
@@ -222,8 +224,8 @@
 
       // Find the target dispatch to break on and trace on all matching
       // dispatches.
-      DispatchOp breakTarget = DispatchOp();
-      funcOp.walk([&](DispatchOp dispatchOp) {
+      IREE::Flow::DispatchOp breakTarget;
+      funcOp.walk([&](IREE::Flow::DispatchOp dispatchOp) {
         std::string entryPointName =
             dispatchOp.getEntryPoint().getRootReference().getValue().str();
         for (FlatSymbolRefAttr nestedRef :
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/inject_dispatch_tracing.mlir b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/inject_dispatch_tracing.mlir
index 7471d55..c090fd6 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/inject_dispatch_tracing.mlir
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/inject_dispatch_tracing.mlir
@@ -4,10 +4,10 @@
 // CHECK-SAME: (%[[ARG0:.+]]: tensor<4xf32>)
 func.func @singleDispatch(%arg0: tensor<4xf32>) -> tensor<4xf32> {
   %c4 = arith.constant 4 : index
-  //      CHECK: flow.tensor.trace {key = "ex::entry0 inputs"} %[[ARG0]] : tensor<4xf32>
+  //      CHECK: flow.tensor.trace "ex::entry0 inputs" = [%[[ARG0]] : tensor<4xf32>]
   // CHECK-NEXT: %[[RET0:.+]] = flow.dispatch @ex::@entry0[%c4](%[[ARG0]]) : (tensor<4xf32>) -> tensor<4xf32>
   %0 = flow.dispatch @ex::@entry0[%c4](%arg0) : (tensor<4xf32>) -> tensor<4xf32>
-  // CHECK-NEXT: flow.tensor.trace {key = "ex::entry0 outputs"} %[[RET0]] : tensor<4xf32>
+  // CHECK-NEXT: flow.tensor.trace "ex::entry0 outputs" = [%[[RET0]] : tensor<4xf32>]
   // CHECK-NEXT: return %[[RET0]]
   return %0 : tensor<4xf32>
 }
@@ -19,15 +19,15 @@
 func.func @multiDispatch(%arg0: tensor<4xf32>) -> tensor<4xf32> {
   %c4 = arith.constant 4 : index
 
-  //     CHECK: flow.tensor.trace {key = "ex::entry0 inputs"} %[[ARG0]] : tensor<4xf32>
+  //      CHECK: flow.tensor.trace "ex::entry0 inputs" = [%[[ARG0]] : tensor<4xf32>]
   // CHECK-NEXT: %[[RET0:.+]] = flow.dispatch @ex::@entry0[%c4](%[[ARG0]]) : (tensor<4xf32>) -> tensor<4xf32>
   %0 = flow.dispatch @ex::@entry0[%c4](%arg0) : (tensor<4xf32>) -> tensor<4xf32>
-  // CHECK-NEXT: flow.tensor.trace {key = "ex::entry0 outputs"} %[[RET0]] : tensor<4xf32>
+  // CHECK-NEXT: flow.tensor.trace "ex::entry0 outputs" = [%[[RET0]] : tensor<4xf32>]
 
-  //     CHECK: flow.tensor.trace {key = "ex::entry1 inputs"} %[[RET0]] : tensor<4xf32>
+  //      CHECK: flow.tensor.trace "ex::entry1 inputs" = [%[[RET0]] : tensor<4xf32>]
   // CHECK-NEXT: %[[RET1:.+]] = flow.dispatch @ex::@entry1[%c4](%[[RET0]]) : (tensor<4xf32>) -> tensor<4xf32>
   %1 = flow.dispatch @ex::@entry1[%c4](%0) : (tensor<4xf32>) -> tensor<4xf32>
-  // CHECK-NEXT: flow.tensor.trace {key = "ex::entry1 outputs"} %[[RET1]] : tensor<4xf32>
+  // CHECK-NEXT: flow.tensor.trace "ex::entry1 outputs" = [%[[RET1]] : tensor<4xf32>]
 
   // CHECK: return %[[RET1]]
   return %1 : tensor<4xf32>
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/insert_dispatch_debug_markers.mlir b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/insert_dispatch_debug_markers.mlir
index f17537b..c15f268 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/insert_dispatch_debug_markers.mlir
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/insert_dispatch_debug_markers.mlir
@@ -2,17 +2,17 @@
 // RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(iree-flow-insert-debug-target-at-ordinal{break-debug-target=@target_func:0 trace-debug-target=@target_func:0})" %s | FileCheck %s --check-prefixes=ORDINAL_0
 // RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(iree-flow-insert-debug-target-at-symbol{break-debug-target=dispatch_1 trace-debug-target=dispatch_1[^0-9]})" %s | FileCheck %s --check-prefixes=CHECK,SYMBOL
 
-/// Multiple functions
+// Multiple functions.
 
 // CHECK-LABEL: func.func @target_func
 // ORDINAL_0-LABEL: func.func @target_func
 func.func @target_func(%arg0: tensor<4xf32>) -> !hal.buffer_view {
   %c4 = arith.constant 4 : index
   // CHECK: %[[D0:.+]] = flow.dispatch @dispatch_0::@dispatch_0_entry
-  //      ORDINAL_0: flow.tensor.trace {key = "dispatch_0::dispatch_0_entry::0 inputs"}
+  //      ORDINAL_0: flow.tensor.trace "dispatch_0::dispatch_0_entry::0 inputs"
   // ORDINAL_0-NEXT: %[[D0:.+]] = flow.dispatch @dispatch_0::@dispatch_0_entry
   %0 = flow.dispatch @dispatch_0::@dispatch_0_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
-  // ORDINAL_0-NEXT: flow.tensor.trace {key = "dispatch_0::dispatch_0_entry::0 outputs"}
+  // ORDINAL_0-NEXT: flow.tensor.trace "dispatch_0::dispatch_0_entry::0 outputs"
   // CHECK: %[[D1:.+]] = flow.dispatch @dispatch_1::@dispatch_1_entry
   %1 = flow.dispatch @dispatch_1::@dispatch_1_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
   %2 = flow.dispatch @dispatch_2::@dispatch_2_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
@@ -22,23 +22,22 @@
   return %3 : !hal.buffer_view
 }
 
-
 // CHECK-LABEL: func.func @other_func
 func.func @other_func(%arg0: tensor<4xf32>) -> !hal.buffer_view {
   %c4 = arith.constant 4 : index
-  // CHECK: %[[D0:.+]] = flow.dispatch @dispatch_1::@dispatch_1_entry
-  %0 = flow.dispatch @dispatch_1::@dispatch_1_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
+  // CHECK: %[[D3:.+]] = flow.dispatch @dispatch_3::@dispatch_3_entry
+  %0 = flow.dispatch @dispatch_3::@dispatch_3_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
 
-  // CHECK: %[[D1:.+]] = flow.dispatch @dispatch_2::@dispatch_2_entry
-  %1 = flow.dispatch @dispatch_2::@dispatch_2_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
-  // CHECK: %[[D2:.+]] = flow.dispatch @dispatch_3::@dispatch_3_entry
-  %2 = flow.dispatch @dispatch_3::@dispatch_3_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
+  // CHECK: %[[D4:.+]] = flow.dispatch @dispatch_4::@dispatch_4_entry
+  %1 = flow.dispatch @dispatch_4::@dispatch_4_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
+  // CHECK: %[[D5:.+]] = flow.dispatch @dispatch_5::@dispatch_5_entry
+  %2 = flow.dispatch @dispatch_5::@dispatch_5_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
 
-  // ORDINAL: %[[ORIGINAL_EXPORT:.+]] = hal.tensor.export %[[D2]] : tensor<4xf32> -> !hal.buffer_view
-  // SYMBOL:  %[[BREAK_EXPORT:.+]] = hal.tensor.export %[[D0]] : tensor<4xf32> -> !hal.buffer_view
+  // ORDINAL: %[[ORIGINAL_EXPORT:.+]] = hal.tensor.export %[[D5]] : tensor<4xf32> -> !hal.buffer_view
+  // SYMBOL:  %[[BREAK_EXPORT:.+]] = hal.tensor.export %[[D5]] : tensor<4xf32> -> !hal.buffer_view
   %3 = hal.tensor.export %2 : tensor<4xf32> -> !hal.buffer_view
 
-  /// Only break on the symbol as the ordinal specifies a different function
+  // Only break on the symbol as the ordinal specifies a different function.
   // SYMBOL:  return %[[BREAK_EXPORT]] : !hal.buffer_view
   // ORDINAL: return %[[ORIGINAL_EXPORT]] : !hal.buffer_view
   return %3 : !hal.buffer_view
@@ -46,7 +45,8 @@
 
 // -----
 
-// Break on a dispatch with a different number of results
+// Break on a dispatch with a different number of results.
+
 // CHECK-LABEL: func.func @target_func
 func.func @target_func(%arg0: tensor<4xf32>) -> !hal.buffer_view {
   %c4 = arith.constant 4 : index
@@ -64,7 +64,8 @@
 
 // -----
 
-// Break/trace on a dispatch not found in the target function should do nothing
+// Break/trace on a dispatch not found in the target function should do nothing.
+
 // CHECK-LABEL: func.func @target_func
 func.func @target_func(%arg0: tensor<4xf32>) -> !hal.buffer_view {
   %c4 = arith.constant 4 : index
@@ -78,7 +79,8 @@
 
 // -----
 
-/// Combine tracing and breaking on the same dispatch
+// Combines tracing and breaking on the same dispatch.
+
 // CHECK-LABEL: func.func @target_func
 // CHECK-SAME:       %[[ARG0:.+]]: tensor<4xf32>
 func.func @target_func(%arg0: tensor<4xf32>) -> !hal.buffer_view {
@@ -86,12 +88,12 @@
   // CHECK: %[[D0:.+]] = flow.dispatch @dispatch_0::@dispatch_0_entry
   %0 = flow.dispatch @dispatch_0::@dispatch_0_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
 
-  // ORDINAL: flow.tensor.trace {key = "dispatch_1::dispatch_1_entry::1 inputs"} %[[ARG0]] : tensor<4xf32>
-  // SYMBOL:  flow.tensor.trace {key = "dispatch_1::dispatch_1_entry inputs"} %[[ARG0]] : tensor<4xf32>
+  // ORDINAL: flow.tensor.trace "dispatch_1::dispatch_1_entry::1 inputs" = [%[[ARG0]] : tensor<4xf32>]
+  // SYMBOL:  flow.tensor.trace "dispatch_1::dispatch_1_entry inputs" = [%[[ARG0]] : tensor<4xf32>]
   // CHECK: %[[D1:.+]] = flow.dispatch @dispatch_1::@dispatch_1_entry
   %1 = flow.dispatch @dispatch_1::@dispatch_1_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
-  // ORDINAL: flow.tensor.trace {key = "dispatch_1::dispatch_1_entry::1 outputs"} %[[D1]] : tensor<4xf32>
-  // SYMBOL:  flow.tensor.trace {key = "dispatch_1::dispatch_1_entry outputs"} %[[D1]] : tensor<4xf32>
+  // ORDINAL: flow.tensor.trace "dispatch_1::dispatch_1_entry::1 outputs" = [%[[D1]] : tensor<4xf32>]
+  // SYMBOL:  flow.tensor.trace "dispatch_1::dispatch_1_entry outputs" = [%[[D1]] : tensor<4xf32>]
 
   %2 = flow.dispatch @dispatch_2::@dispatch_2_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
   %3 = hal.tensor.export %2 : tensor<4xf32> -> !hal.buffer_view
@@ -103,19 +105,21 @@
 
 // -----
 
-/// Check regex matching on symbol
+// Checks regex matching on a dispatch symbol.
+
 // CHECK-LABEL: func.func @target_func
 func.func @target_func(%arg0: tensor<4xf32>) -> !hal.buffer_view {
   %c4 = arith.constant 4 : index
-  // SYMBOL: flow.tensor.trace {key = "dispatch_1::dispatch_1_entry inputs"}
+
+  // SYMBOL: flow.tensor.trace "dispatch_1::dispatch_1_entry inputs"
   // CHECK:  flow.dispatch @dispatch_1::@dispatch_1_entry
   %0 = flow.dispatch @dispatch_1::@dispatch_1_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
-  // SYMBOL: flow.tensor.trace {key = "dispatch_1::dispatch_1_entry outputs"}
+  // SYMBOL: flow.tensor.trace "dispatch_1::dispatch_1_entry outputs"
 
-  // SYMBOL-NOT: flow.tensor.trace {key = "dispatch_11::dispatch_11_entry inputs"}
+  // SYMBOL-NOT: flow.tensor.trace "dispatch_11::dispatch_11_entry inputs"
   // CHECK:      flow.dispatch @dispatch_11::@dispatch_11_entry
   %1 = flow.dispatch @dispatch_11::@dispatch_11_entry[%c4] (%arg0) : (tensor<4xf32>) -> tensor<4xf32>
-  // SYMBOL-NOT: flow.tensor.trace {key = "dispatch_11::dispatch_11_entry outputs"}
+  // SYMBOL-NOT: flow.tensor.trace "dispatch_11::dispatch_11_entry outputs"
 
   %2 = hal.tensor.export %1 : tensor<4xf32> -> !hal.buffer_view
   return %2 : !hal.buffer_view
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp
index 8d31a02..9c7ebe0 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp
@@ -758,8 +758,22 @@
   LogicalResult
   matchAndRewrite(IREE::Stream::TensorTraceOp traceOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
+    SmallVector<Value> bufferViews;
+    auto resourceEncodingDims = adaptor.getResourceEncodingDims();
+    for (auto [resource, resourceSize, resourceEncoding] : llvm::zip_equal(
+             adaptor.getResources(), adaptor.getResourceSizes(),
+             adaptor.getResourceEncodings().getAsRange<TypeAttr>())) {
+      int64_t dynamicDimCount =
+          cast<ShapedType>(resourceEncoding.getValue()).getNumDynamicDims();
+      bufferViews.push_back(rewriter.create<IREE::Stream::TensorExportOp>(
+          traceOp.getLoc(), rewriter.getType<IREE::HAL::BufferViewType>(),
+          resource, resourceEncoding,
+          resourceEncodingDims.take_front(dynamicDimCount), resourceSize,
+          /*affinity=*/IREE::Stream::AffinityAttr{}));
+      resourceEncodingDims = resourceEncodingDims.drop_front(dynamicDimCount);
+    }
     rewriter.replaceOpWithNewOp<IREE::HAL::BufferViewTraceOp>(
-        traceOp, traceOp.getKeyAttr(), adaptor.getOperands());
+        traceOp, traceOp.getKeyAttr(), bufferViews);
     return success();
   }
 };
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/BUILD.bazel b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/BUILD.bazel
index 832da44..64493c0 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/BUILD.bazel
@@ -18,6 +18,7 @@
         [
             "channel_ops.mlir",
             "cmd_ops.mlir",
+            "debug_ops.mlir",
             "file_ops.mlir",
             "resource_ops.mlir",
             "timepoint_ops.mlir",
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/CMakeLists.txt
index 5c25a89..fb7c40f 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/CMakeLists.txt
@@ -16,6 +16,7 @@
   SRCS
     "channel_ops.mlir"
     "cmd_ops.mlir"
+    "debug_ops.mlir"
     "file_ops.mlir"
     "resource_ops.mlir"
     "timepoint_ops.mlir"
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/debug_ops.mlir b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/debug_ops.mlir
new file mode 100644
index 0000000..a748326
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/debug_ops.mlir
@@ -0,0 +1,14 @@
+// RUN: iree-opt --split-input-file --iree-hal-conversion %s | FileCheck %s
+
+// CHECK-LABEL: @tensorTrace
+// CHECK-SAME: (%[[TENSOR0_BUFFER:.+]]: !hal.buffer, %[[TENSOR0_SIZE:.+]]: index, %[[TENSOR1_BUFFER:.+]]: !hal.buffer, %[[TENSOR1_SIZE:.+]]: index, %[[TENSOR1_DIM0:.+]]: index)
+func.func @tensorTrace(%tensor0: !stream.resource<staging>, %tensor0_size: index, %tensor1: !stream.resource<staging>, %tensor1_size: index, %tensor1_dim0: index) {
+  // CHECK-DAG: %[[TENSOR0:.+]] = hal.buffer_view.create buffer(%[[TENSOR0_BUFFER]] : !hal.buffer)[%c0{{.*}}, %[[TENSOR0_SIZE]]] shape([%c5, %c3])
+  // CHECK-DAG: %[[TENSOR1:.+]] =  hal.buffer_view.create buffer(%[[TENSOR1_BUFFER]] : !hal.buffer)[%c0{{.*}}, %[[TENSOR1_SIZE]]] shape([%[[TENSOR1_DIM0]], %c5{{.*}}])
+  // CHECK: hal.buffer_view.trace "FOOBAR" = %[[TENSOR0]], %[[TENSOR1]] : !hal.buffer_view, !hal.buffer_view
+  stream.tensor.trace "FOOBAR" = [
+    %tensor0 : tensor<5x3xf32> in !stream.resource<staging>{%tensor0_size},
+    %tensor1 : tensor<?x5xf32>{%tensor1_dim0} in !stream.resource<staging>{%tensor1_size}
+  ]
+  return
+}
diff --git a/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td b/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td
index bc58f6f..e2f795e 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td
+++ b/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td
@@ -880,6 +880,7 @@
   );
 
   let assemblyFormat = [{
+    $key `=`
     $operands `:` type($operands)
     attr-dict-with-keyword
   }];
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/ResourceUsage.cpp b/compiler/src/iree/compiler/Dialect/Stream/Analysis/ResourceUsage.cpp
index e7e2bed..e036755 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Analysis/ResourceUsage.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Analysis/ResourceUsage.cpp
@@ -591,6 +591,9 @@
         .Case([&](IREE::Stream::TensorExportOp op) {
           removeAssumedBits(NOT_MUTATED | NOT_EXTERNAL);
         })
+        .Case([&](IREE::Stream::TensorTraceOp op) {
+          removeAssumedBits(NOT_STAGING_READ);
+        })
         .Case([&](IREE::Stream::AsyncCloneOp op) {
           removeAssumedBits(NOT_TRANSFER_READ);
           auto &resultUsage = solver.getElementFor<ValueResourceUsage>(
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp
index 297eba6..4ae1571 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp
@@ -244,30 +244,30 @@
   LogicalResult
   matchAndRewrite(IREE::Flow::TensorTraceOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    SmallVector<Value> exportedTensors;
+    SmallVector<Value> resources;
+    SmallVector<Value> resourceSizes;
+    SmallVector<Attribute> resourceEncodings;
     for (auto [tensorOperand, resourceOperand] :
-         llvm::zip_equal(op.getOperands(), adaptor.getOperands())) {
+         llvm::zip_equal(op.getValues(), adaptor.getValues())) {
       auto source =
           consumeTensorOperand(op.getLoc(), resourceOperand, rewriter);
-      auto externalType = rewriter.getType<IREE::Stream::ResourceType>(
-          IREE::Stream::Lifetime::External);
-      auto exportSource = resourceOperand;
-      if (source.resource.getType() != externalType) {
-        exportSource = rewriter.create<IREE::Stream::AsyncTransferOp>(
-            op.getLoc(), externalType, source.resource, source.resourceSize,
+      auto stagingType = rewriter.getType<IREE::Stream::ResourceType>(
+          IREE::Stream::Lifetime::Staging);
+      auto traceSource = source.resource;
+      if (source.resource.getType() != stagingType) {
+        traceSource = rewriter.create<IREE::Stream::AsyncTransferOp>(
+            op.getLoc(), stagingType, source.resource, source.resourceSize,
             source.resourceSize,
             /*source_affinity=*/getAffinityFor(op),
             /*result_affinity=*/nullptr);
       }
-      auto dynamicDims = IREE::Util::buildDynamicDimsForValue(
-          op.getLoc(), tensorOperand, rewriter);
-      exportedTensors.push_back(rewriter.create<IREE::Stream::TensorExportOp>(
-          op.getLoc(), tensorOperand.getType(), exportSource,
-          TypeAttr::get(tensorOperand.getType()), dynamicDims,
-          source.resourceSize, getAffinityFor(op)));
+      resources.push_back(traceSource);
+      resourceSizes.push_back(source.resourceSize);
+      resourceEncodings.push_back(TypeAttr::get(tensorOperand.getType()));
     }
     rewriter.replaceOpWithNewOp<IREE::Stream::TensorTraceOp>(
-        op, adaptor.getKey(), exportedTensors);
+        op, adaptor.getKey(), resources, resourceSizes,
+        rewriter.getArrayAttr(resourceEncodings), adaptor.getValueDims());
     return success();
   }
 };
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/tensor_ops.mlir b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/tensor_ops.mlir
index d194376..4e2a972 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/tensor_ops.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/tensor_ops.mlir
@@ -147,3 +147,21 @@
   // CHECK: return %[[T2]]
   return %0 : tensor<2x3xi32>
 }
+
+// -----
+
+// CHECK-LABEL: @tensorTrace
+//  CHECK-SAME: (%[[TENSOR0:.+]]: !stream.resource<*>, %[[TENSOR0_SIZE:.+]]: index, %[[TENSOR1:.+]]: !stream.resource<*>, %[[TENSOR1_SIZE:.+]]: index, %[[TENSOR1_DIM0:.+]]: index, %[[TENSOR1_DIM2:.+]]: index)
+func.func @tensorTrace(%tensor0: tensor<5xf32>, %tensor1: tensor<?x3x?xi32>, %tensor1_dim0: index, %tensor1_dim2: index) {
+  // CHECK-DAG: %[[TENSOR0_STAGED:.+]] = stream.async.transfer %[[TENSOR0]] : !stream.resource<*>{%[[TENSOR0_SIZE]]} -> !stream.resource<staging>{%[[TENSOR0_SIZE]]}
+  // CHECK-DAG: %[[TENSOR1_STAGED:.+]] = stream.async.transfer %[[TENSOR1]] : !stream.resource<*>{%[[TENSOR1_SIZE]]} -> !stream.resource<staging>{%[[TENSOR1_SIZE]]}
+  //      CHECK: stream.tensor.trace "FOOBAR" = [
+  // CHECK-NEXT:   %[[TENSOR0_STAGED]] : tensor<5xf32> in !stream.resource<staging>{%[[TENSOR0_SIZE]]},
+  // CHECK-NEXT:   %[[TENSOR1_STAGED]] : tensor<?x3x?xi32>{%[[TENSOR1_DIM0]], %[[TENSOR1_DIM2]]} in !stream.resource<staging>{%[[TENSOR1_SIZE]]}
+  // CHECK-NEXT: ]
+  flow.tensor.trace "FOOBAR" = [
+    %tensor0 : tensor<5xf32>,
+    %tensor1 : tensor<?x3x?xi32>{%tensor1_dim0, %tensor1_dim2}
+  ]
+  return
+}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
index addada8..a294923 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
@@ -192,6 +192,20 @@
   return success();
 }
 
+static void eraseStreamRegionResults(Region &region,
+                                     ArrayRef<unsigned> excludedResultIndices) {
+  for (auto &block : region.getBlocks()) {
+    auto yieldOp = dyn_cast<IREE::Stream::YieldOp>(block.getTerminator());
+    if (!yieldOp)
+      continue;
+    llvm::SmallVector<Value> newOperands;
+    for (auto i : llvm::reverse(excludedResultIndices)) {
+      yieldOp.getResourceOperandsMutable().erase(i);
+      yieldOp.getResourceOperandSizesMutable().erase(i);
+    }
+  }
+}
+
 // Computes the value access bits starting from |rootValue|.
 // Traverses the IR graph along tied ops but does not handle branches.
 static IREE::Util::ValueAccess computeValueAccess(Value rootValue) {
@@ -260,18 +274,82 @@
   return access;
 }
 
-static void eraseStreamRegionResults(Region &region,
-                                     ArrayRef<unsigned> excludedResultIndices) {
-  for (auto &block : region.getBlocks()) {
-    auto yieldOp = dyn_cast<IREE::Stream::YieldOp>(block.getTerminator());
-    if (!yieldOp)
-      continue;
-    llvm::SmallVector<Value> newOperands;
-    for (auto i : llvm::reverse(excludedResultIndices)) {
-      yieldOp.getResourceOperandsMutable().erase(i);
-      yieldOp.getResourceOperandSizesMutable().erase(i);
+//===----------------------------------------------------------------------===//
+// custom<EncodedResourceOperands>(
+//     $resources, type($resources), $resource_sizes,
+//     $resource_encodings, $resource_encoding_dims)
+//===----------------------------------------------------------------------===//
+
+static ParseResult parseEncodedResourceOperands(
+    OpAsmParser &parser,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resources,
+    SmallVectorImpl<Type> &resourceTypes,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resourceSizes,
+    ArrayAttr &resourceEncodings,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resourceEncodingDims) {
+  SmallVector<Attribute> resourceEncodingAttrs;
+  do {
+    resources.emplace_back();
+    TypeAttr resourceEncoding;
+    if (failed(parser.parseOperand(resources.back())) ||
+        failed(parser.parseColon()) ||
+        failed(parser.parseAttribute(resourceEncoding)))
+      return failure();
+    resourceEncodingAttrs.push_back(resourceEncoding);
+    if (int64_t dynamicDimCount =
+            cast<ShapedType>(resourceEncoding.getValue()).getNumDynamicDims()) {
+      if (failed(parser.parseOperandList(resourceEncodingDims, dynamicDimCount,
+                                         AsmParser::Delimiter::Braces)))
+        return failure();
     }
-  }
+    resourceTypes.emplace_back();
+    resourceSizes.emplace_back();
+    if (failed(parser.parseKeyword("in")) ||
+        failed(parseSizeAwareType(parser, resourceTypes.back(),
+                                  resourceSizes.back())))
+      return failure();
+  } while (succeeded(parser.parseOptionalComma()));
+  resourceEncodings = parser.getBuilder().getArrayAttr(resourceEncodingAttrs);
+  return success();
+}
+
+static void printEncodedResourceOperands(OpAsmPrinter &p, Operation *op,
+                                         ValueRange resources,
+                                         TypeRange resourceTypes,
+                                         ValueRange resourceSizes,
+                                         ArrayAttr resourceEncodings,
+                                         ValueRange resourceEncodingDims) {
+  p.increaseIndent();
+  p.printNewline();
+  llvm::interleave(
+      llvm::zip_equal(resources, resourceTypes, resourceSizes,
+                      resourceEncodings.getAsValueRange<TypeAttr>()),
+      [&](auto it) {
+        auto [resource, resourceType, resourceSize, resourceEncoding] = it;
+        p << resource;
+        p << " : ";
+        p << resourceEncoding;
+        if (int64_t dynamicDimCount =
+                cast<ShapedType>(resourceEncoding).getNumDynamicDims()) {
+          p << "{";
+          llvm::interleaveComma(
+              resourceEncodingDims.take_front(dynamicDimCount), p);
+          resourceEncodingDims =
+              resourceEncodingDims.drop_front(dynamicDimCount);
+          p << "}";
+        }
+        p << " in ";
+        p << resourceType;
+        p << "{";
+        p << resourceSize;
+        p << "}";
+      },
+      [&]() {
+        p << ",";
+        p.printNewline();
+      });
+  p.decreaseIndent();
+  p.printNewline();
 }
 
 //===----------------------------------------------------------------------===//
@@ -1212,6 +1290,54 @@
 }
 
 //===----------------------------------------------------------------------===//
+// stream.tensor.trace
+//===----------------------------------------------------------------------===//
+
+LogicalResult TensorTraceOp::verify() {
+  TensorTraceOp op = *this;
+  if (op.getResources().size() != op.getResourceEncodings().size() ||
+      op.getResources().size() != op.getResourceSizes().size()) {
+    return op.emitOpError(
+        "each resource needs a matching resource encoding and size "
+        "(array length mismatch)");
+  }
+  auto resourceEncodingDims = op.getResourceEncodingDims();
+  for (auto [resource, resourceSize, resourceEncoding] :
+       llvm::zip_equal(op.getResources(), op.getResourceSizes(),
+                       op.getResourceEncodings().getAsValueRange<TypeAttr>())) {
+    int64_t dynamicDimCount =
+        cast<ShapedType>(resourceEncoding).getNumDynamicDims();
+    if (failed(verifyOpDynamicDims(
+            op, resourceEncoding,
+            resourceEncodingDims.take_front(dynamicDimCount))) ||
+        failed(verifyOpValueSizes(op, resource, resourceSize))) {
+      return failure();
+    }
+    resourceEncodingDims = resourceEncodingDims.drop_front(dynamicDimCount);
+  }
+  return success();
+}
+
+ValueRange TensorTraceOp::getOperandDynamicDims(unsigned idx) {
+  auto resourceEncodings = getResourceEncodings().getValue();
+  auto resourceEncodingDims = getResourceEncodingDims();
+  for (unsigned i = 0; i <= idx; ++i) {
+    auto resourceEncoding = resourceEncodings[i].cast<TypeAttr>().getValue();
+    int64_t dynamicDimCount =
+        cast<ShapedType>(resourceEncoding).getNumDynamicDims();
+    if (i == idx) {
+      return resourceEncodingDims.take_front(dynamicDimCount);
+    }
+    resourceEncodingDims = resourceEncodingDims.drop_front(dynamicDimCount);
+  }
+  return ValueRange{};
+}
+
+ValueRange TensorTraceOp::getResultDynamicDims(unsigned idx) {
+  return ValueRange{};
+}
+
+//===----------------------------------------------------------------------===//
 // stream.async.alloca
 //===----------------------------------------------------------------------===//
 
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
index 275c178..0d52718 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
@@ -1383,22 +1383,40 @@
   let hasCanonicalizer = 1;
 }
 
-def Stream_TensorTraceOp : Stream_Op<"tensor.trace", []> {
-  let summary = [{trace value(s) operation}];
+def Stream_TensorTraceOp : Stream_Op<"tensor.trace", [
+  AttrSizedOperandSegments,
+  DeclareOpInterfaceMethods<Util_ShapeAwareOp>,
+  Util_SizeAwareOp,
+]> {
+  let summary = [{traces one or more tensor values at runtime}];
   let description = [{
     Traces out to a runtime trace sink (console, log file, etc) the given
-    tensors and titles them with the given key. The key is informational only
-    and useful for titling/marking specific sets of tensors for easier
-    searching.
+    tensors. The key is arbitrary and can be used for identifying the set of
+    values being traced.
   }];
 
   let arguments = (ins
     StrAttr:$key,
-    Variadic<AnyTensor>:$operands
+    Variadic<Stream_StagingResource>:$resources,
+    Variadic<Stream_Size>:$resource_sizes,
+    TypeArrayAttr:$resource_encodings,
+    Stream_ShapeDynamicDims:$resource_encoding_dims
   );
 
-  // TODO(benvanik): stream.tensor.trace assembly format (custom).
-  let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
+  let assemblyFormat = [{
+    $key `=` `[`
+    custom<EncodedResourceOperands>(
+        $resources, type($resources), $resource_sizes,
+        $resource_encodings, $resource_encoding_dims)
+    `]` attr-dict-with-keyword
+  }];
+
+  let extraClassDeclaration = [{
+    Value getOperandSize(unsigned idx) { return getResourceSizes()[idx]; }
+    Value getResultSize(unsigned idx) { return {}; }
+  }];
+
+  let hasVerifier = 1;
 }
 
 } // OpGroupTensorOps
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/test/tensor_ops.mlir b/compiler/src/iree/compiler/Dialect/Stream/IR/test/tensor_ops.mlir
index f650415..4cc586c 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/test/tensor_ops.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/test/tensor_ops.mlir
@@ -135,3 +135,19 @@
   %0 = stream.tensor.store %arg2, %arg0 : f32 -> tensor<f32> in %arg0 as !stream.resource<staging>{%arg1}
   return %0 : !stream.resource<staging>
 }
+
+// -----
+
+// CHECK-LABEL: @tensorTrace
+//  CHECK-SAME: (%[[TENSOR0:.+]]: !stream.resource<staging>, %[[TENSOR0_SIZE:.+]]: index, %[[TENSOR1:.+]]: !stream.resource<staging>, %[[TENSOR1_SIZE:.+]]: index, %[[TENSOR1_DIM0:.+]]: index, %[[TENSOR1_DIM2:.+]]: index)
+func.func @tensorTrace(%tensor0: !stream.resource<staging>, %tensor0_size: index, %tensor1: !stream.resource<staging>, %tensor1_size: index, %tensor1_dim0: index, %tensor1_dim2: index) {
+  //      CHECK: stream.tensor.trace "FOOBAR" = [
+  // CHECK-NEXT:   %[[TENSOR0]] : tensor<5xf32> in !stream.resource<staging>{%[[TENSOR0_SIZE]]},
+  // CHECK-NEXT:   %[[TENSOR1]] : tensor<?x3x?xi32>{%[[TENSOR1_DIM0]], %[[TENSOR1_DIM2]]} in !stream.resource<staging>{%[[TENSOR1_SIZE]]}
+  // CHECK-NEXT: ]
+  stream.tensor.trace "FOOBAR" = [
+    %tensor0 : tensor<5xf32> in !stream.resource<staging>{%tensor0_size},
+    %tensor1 : tensor<?x3x?xi32>{%tensor1_dim0, %tensor1_dim2} in !stream.resource<staging>{%tensor1_size}
+  ]
+  return
+}
diff --git a/compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.cpp b/compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.cpp
index b13b245..b8cbc9c 100644
--- a/compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.cpp
+++ b/compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.cpp
@@ -727,6 +727,15 @@
   return dynamicDims;
 }
 
+SmallVector<Value> buildDynamicDimsForValues(Location loc, ValueRange values,
+                                             OpBuilder &builder) {
+  SmallVector<Value> dynamicDims;
+  for (auto value : values) {
+    dynamicDims.append(buildDynamicDimsForValue(loc, value, builder));
+  }
+  return dynamicDims;
+}
+
 static SmallVector<Value> buildShape(Location loc, ShapedType type,
                                      ValueRange dynamicDims,
                                      OpBuilder &builder) {
diff --git a/compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.h b/compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.h
index 76bb1c1..f27ee11 100644
--- a/compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.h
+++ b/compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.h
@@ -170,6 +170,12 @@
 SmallVector<Value> buildDynamicDimsForValue(Location loc, Value value,
                                             OpBuilder &builder);
 
+// Returns dimension values for each dynamic dimension of the given |values|.
+// |values| must all have ShapedTypes. The returned value range will be empty if
+// all shapes are fully static.
+SmallVector<Value> buildDynamicDimsForValues(Location loc, ValueRange values,
+                                             OpBuilder &builder);
+
 // Builds a ranked shape with all dimension values for the given operand.
 SmallVector<Value> buildOperandShape(ShapeAwareOpInterface op,
                                      unsigned operandIdx, OpBuilder &builder);
diff --git a/compiler/src/iree/compiler/InputConversion/Common/test/iree_import_public.mlir b/compiler/src/iree/compiler/InputConversion/Common/test/iree_import_public.mlir
index f9b3730..5aa7bbf 100644
--- a/compiler/src/iree/compiler/InputConversion/Common/test/iree_import_public.mlir
+++ b/compiler/src/iree/compiler/InputConversion/Common/test/iree_import_public.mlir
@@ -62,7 +62,7 @@
 // CHECK: %buffer = hal.buffer.subspan
 // CHECK-SAME: <%arg0 : !hal.buffer>[%[[OFFSET]], %[[LENGTH]]] : !hal.buffer
 
-func.func @buffer_subspan(%arg0: !iree_input.buffer) -> !iree_input.buffer {  
+func.func @buffer_subspan(%arg0: !iree_input.buffer) -> !iree_input.buffer {
   %offset = arith.constant 100 : index
   %length = arith.constant 200 : index
   %buffer = iree_input.buffer.subspan<%arg0 : !iree_input.buffer>[%offset, %length] : !iree_input.buffer
@@ -275,9 +275,15 @@
 
 // -----
 // CHECK-LABEL: func.func @tensor_trace
-// CHECK: flow.tensor.trace {key = "FOOBAR"} %arg0, %arg1 : tensor<5xf32>, tensor<3xf32>
-func.func @tensor_trace(%arg0 : tensor<5xf32>, %arg1 : tensor<3xf32>) {
-  iree_input.tensor.trace "FOOBAR" %arg0, %arg1 : tensor<5xf32>, tensor<3xf32>
+//      CHECK: flow.tensor.trace "FOOBAR" = [
+// CHECK-SAME:   %arg0 : tensor<5xf32>,
+// CHECK-SAME:   %arg1 : tensor<?x3xf32>{%arg2}
+// CHECK-SAME: ]
+func.func @tensor_trace(%arg0: tensor<5xf32>, %arg1: tensor<?x3xf32>, %arg2: index) {
+  iree_input.tensor.trace "FOOBAR" = [
+    %arg0 : tensor<5xf32>,
+    %arg1 : tensor<?x3xf32>{%arg2}
+  ]
   return
 }
 
diff --git a/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/Patterns.cpp b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/Patterns.cpp
index fda46c4..f21e24c 100644
--- a/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/Patterns.cpp
+++ b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/Patterns.cpp
@@ -374,8 +374,29 @@
   LogicalResult
   matchAndRewrite(IREE::Stream::TensorTraceOp traceOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
+    auto bufferType = rewriter.getType<IREE::HAL::BufferType>();
+    auto bufferViewType = rewriter.getType<IREE::HAL::BufferViewType>();
+    auto zero = rewriter.create<arith::ConstantIndexOp>(traceOp.getLoc(), 0);
+    auto resourceEncodingDims = adaptor.getResourceEncodingDims();
+    SmallVector<Value> bufferViews;
+    for (auto [resource, resourceSize, resourceEncoding] : llvm::zip_equal(
+             adaptor.getResources(), adaptor.getResourceSizes(),
+             adaptor.getResourceEncodings().getAsRange<TypeAttr>())) {
+      Value resourceBuffer = rewriter.create<IREE::HAL::Inline::BufferWrapOp>(
+          traceOp.getLoc(), bufferType, resource,
+          /*offset=*/
+          zero,
+          /*length=*/resourceSize);
+      int64_t dynamicDimCount =
+          cast<ShapedType>(resourceEncoding.getValue()).getNumDynamicDims();
+      bufferViews.push_back(rewriter.create<IREE::Stream::TensorExportOp>(
+          traceOp.getLoc(), bufferViewType, resourceBuffer, resourceEncoding,
+          resourceEncodingDims.take_front(dynamicDimCount), resourceSize,
+          /*affinity=*/IREE::Stream::AffinityAttr{}));
+      resourceEncodingDims = resourceEncodingDims.drop_front(dynamicDimCount);
+    }
     rewriter.replaceOpWithNewOp<IREE::HAL::Inline::BufferViewTraceOp>(
-        traceOp, traceOp.getKeyAttr(), adaptor.getOperands());
+        traceOp, traceOp.getKeyAttr(), bufferViews);
     return success();
   }
 };
diff --git a/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/BUILD.bazel b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/BUILD.bazel
index db19be9..09466b8 100644
--- a/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/BUILD.bazel
@@ -17,6 +17,7 @@
     srcs = enforce_glob(
         [
             "cmd_ops.mlir",
+            "debug_ops.mlir",
             "file_ops.mlir",
             "resource_ops.mlir",
             "timepoint_ops.mlir",
diff --git a/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/CMakeLists.txt b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/CMakeLists.txt
index 930427f..c3f6ba0 100644
--- a/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/CMakeLists.txt
@@ -15,6 +15,7 @@
     lit
   SRCS
     "cmd_ops.mlir"
+    "debug_ops.mlir"
     "file_ops.mlir"
     "resource_ops.mlir"
     "timepoint_ops.mlir"
diff --git a/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/cmd_ops.mlir b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/cmd_ops.mlir
index 72549c1..cf6ff97 100644
--- a/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/cmd_ops.mlir
+++ b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/cmd_ops.mlir
@@ -156,24 +156,3 @@
   } => !stream.timepoint
   return %timepoint : !stream.timepoint
 }
-
-// -----
-
-// CHECK-LABEL: @trace_tensor
-// CHECK-SAME: %[[ARG0:.+]]: !hal.buffer
-func.func @trace_tensor(%arg0: !stream.resource<external>) -> () {
-  // CHECK-DAG: %[[C180:.+]] = arith.constant 180 : index
-  %c180 = arith.constant 180 : index
-
-  // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
-  // CHECK-DAG: %[[C45:.+]] = arith.constant 45 : index
-  // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
-  // CHECK-DAG: %[[C268435488:.+]] = arith.constant 268435488 : i32
-  // CHECK-DAG: %[[C1_i32:.+]] = arith.constant 1 : i32
-  // CHECK: %[[VIEW:.+]] = hal_inline.buffer_view.create buffer(%[[ARG0]] : !hal.buffer)[%[[C0]], %[[C180]]] shape([%[[C1]], %[[C45]]]) type(%[[C268435488]]) encoding(%[[C1_i32]]) : !hal.buffer_view
-  %tensor = stream.tensor.export %arg0 : tensor<1x45xi32> in !stream.resource<external>{%c180} -> tensor<1x45xi32>
-
-  // CHECK: hal_inline.buffer_view.trace %[[VIEW]] : !hal.buffer_view attributes {key = "whatevs"}
-  stream.tensor.trace {key = "whatevs"} %tensor : tensor<1x45xi32>
-  return
-}
diff --git a/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/debug_ops.mlir b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/debug_ops.mlir
new file mode 100644
index 0000000..24d3565
--- /dev/null
+++ b/compiler/src/iree/compiler/Modules/HAL/Inline/Conversion/StreamToHALInline/test/debug_ops.mlir
@@ -0,0 +1,16 @@
+// RUN: iree-opt --split-input-file --iree-hal-inline-conversion %s | FileCheck %s
+
+// CHECK-LABEL: @tensorTrace
+// CHECK-SAME: (%[[TENSOR0_STORAGE:.+]]: !util.buffer, %[[TENSOR0_SIZE:.+]]: index, %[[TENSOR1_STORAGE:.+]]: !util.buffer, %[[TENSOR1_SIZE:.+]]: index, %[[TENSOR1_DIM0:.+]]: index)
+func.func @tensorTrace(%tensor0: !stream.resource<staging>, %tensor0_size: index, %tensor1: !stream.resource<staging>, %tensor1_size: index, %tensor1_dim0: index) {
+  // CHECK-DAG: %[[TENSOR0_BUFFER:.+]] = hal_inline.buffer.wrap source(%[[TENSOR0_STORAGE]] : !util.buffer)[%c0, %[[TENSOR0_SIZE]]] : !hal.buffer
+  // CHECK-DAG: %[[TENSOR0:.+]] = hal_inline.buffer_view.create buffer(%[[TENSOR0_BUFFER]] : !hal.buffer)[%c0{{.*}}, %[[TENSOR0_SIZE]]] shape([%c5, %c3])
+  // CHECK-DAG: %[[TENSOR1_BUFFER:.+]] = hal_inline.buffer.wrap source(%[[TENSOR1_STORAGE]] : !util.buffer)[%c0, %[[TENSOR1_SIZE]]] : !hal.buffer
+  // CHECK-DAG: %[[TENSOR1:.+]] =  hal_inline.buffer_view.create buffer(%[[TENSOR1_BUFFER]] : !hal.buffer)[%c0{{.*}}, %[[TENSOR1_SIZE]]] shape([%[[TENSOR1_DIM0]], %c5{{.*}}])
+  // CHECK: hal_inline.buffer_view.trace "FOOBAR" = %[[TENSOR0]], %[[TENSOR1]] : !hal.buffer_view, !hal.buffer_view
+  stream.tensor.trace "FOOBAR" = [
+    %tensor0 : tensor<5x3xf32> in !stream.resource<staging>{%tensor0_size},
+    %tensor1 : tensor<?x5xf32>{%tensor1_dim0} in !stream.resource<staging>{%tensor1_size}
+  ]
+  return
+}
diff --git a/compiler/src/iree/compiler/Modules/HAL/Inline/IR/HALInlineOps.td b/compiler/src/iree/compiler/Modules/HAL/Inline/IR/HALInlineOps.td
index 4587058..c93d1c8 100644
--- a/compiler/src/iree/compiler/Modules/HAL/Inline/IR/HALInlineOps.td
+++ b/compiler/src/iree/compiler/Modules/HAL/Inline/IR/HALInlineOps.td
@@ -423,6 +423,7 @@
   );
 
   let assemblyFormat = [{
+    $key `=`
     $operands `:` type($operands)
     attr-dict-with-keyword
   }];
diff --git a/compiler/src/iree/compiler/Modules/HAL/Loader/Conversion/StreamToHALLoader/test/cmd_ops.mlir b/compiler/src/iree/compiler/Modules/HAL/Loader/Conversion/StreamToHALLoader/test/cmd_ops.mlir
index 3bde0c6..30d5fa4 100644
--- a/compiler/src/iree/compiler/Modules/HAL/Loader/Conversion/StreamToHALLoader/test/cmd_ops.mlir
+++ b/compiler/src/iree/compiler/Modules/HAL/Loader/Conversion/StreamToHALLoader/test/cmd_ops.mlir
@@ -89,24 +89,3 @@
   // CHECK: return %c0
   return %fence : !stream.timepoint
 }
-
-// -----
-
-// CHECK-LABEL: @trace_tensor
-// CHECK-SAME: %[[ARG0:.+]]: !hal.buffer
-func.func @trace_tensor(%arg0: !stream.resource<external>) -> () {
-  // CHECK-DAG: %[[C180:.+]] = arith.constant 180 : index
-  %c180 = arith.constant 180 : index
-
-  // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
-  // CHECK-DAG: %[[C45:.+]] = arith.constant 45 : index
-  // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
-  // CHECK-DAG: %[[C268435488:.+]] = arith.constant 268435488 : i32
-  // CHECK-DAG: %[[C1_i32:.+]] = arith.constant 1 : i32
-  // CHECK-DAG: %[[VIEW:.+]] = hal_inline.buffer_view.create buffer(%[[ARG0]] : !hal.buffer)[%[[C0]], %[[C180]]] shape([%[[C1]], %[[C45]]]) type(%[[C268435488]]) encoding(%[[C1_i32]]) : !hal.buffer_view
-  %tensor = stream.tensor.export %arg0 : tensor<1x45xi32> in !stream.resource<external>{%c180} -> tensor<1x45xi32>
-
-  // CHECK: hal_inline.buffer_view.trace %[[VIEW]] : !hal.buffer_view attributes {key = "whatevs"}
-  stream.tensor.trace {key = "whatevs"} %tensor : tensor<1x45xi32>
-  return
-}
diff --git a/llvm-external-projects/iree-dialects/BUILD.bazel b/llvm-external-projects/iree-dialects/BUILD.bazel
index 8b8c2d1..408b265 100644
--- a/llvm-external-projects/iree-dialects/BUILD.bazel
+++ b/llvm-external-projects/iree-dialects/BUILD.bazel
@@ -142,6 +142,7 @@
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:InferTypeOpInterface",
         "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
     ],
 )
 
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/Input/InputOps.td b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/Input/InputOps.td
index a78a94b..a7053c6 100644
--- a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/Input/InputOps.td
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/Input/InputOps.td
@@ -681,21 +681,31 @@
   ];
 }
 
-def IREEInput_TensorTraceOp : IREEInput_Op<"tensor.trace", []> {
-  let summary = [{trace value(s) operation}];
+def IREEInput_TensorTraceOp : IREEInput_Op<"tensor.trace", [
+  AttrSizedOperandSegments,
+]> {
+  let summary = [{traces one or more tensor values at runtime}];
   let description = [{
     Traces out to a runtime trace sink (console, log file, etc) the given
-    tensors and titles them with the given key. The key is informational only
-    and useful for titling/marking specific sets of tensors for easier
-    searching.
+    tensors. The key is arbitrary and can be used for identifying the set of
+    values being traced.
   }];
 
   let arguments = (ins
     StrAttr:$key,
-    Variadic<IREEInput_Tensor>:$operands
+    Variadic<IREEInput_Tensor>:$values,
+    IREEInput_ShapeDynamicDims:$value_dims
   );
 
-  let assemblyFormat = "$key attr-dict ($operands^ `:` type($operands))?";
+  let assemblyFormat = [{
+    $key `=` `[`
+    custom<ShapedOperandList>($values, type($values), $value_dims)
+    `]` attr-dict-with-keyword
+  }];
+
+  let builders = [
+    OpBuilder<(ins "StringRef":$key, "ValueRange":$values)>,
+  ];
 }
 
 } // OpGroupTensorOps
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/Input/CMakeLists.txt b/llvm-external-projects/iree-dialects/lib/Dialect/Input/CMakeLists.txt
index bd5df20..31fb4bf 100644
--- a/llvm-external-projects/iree-dialects/lib/Dialect/Input/CMakeLists.txt
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/Input/CMakeLists.txt
@@ -14,6 +14,7 @@
   MLIRIR
   MLIRInferTypeOpInterface
   MLIRSideEffectInterfaces
+  MLIRTensorDialect
 )
 
 iree_dialects_target_includes(IREEInputDialect)
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/Input/InputOps.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/Input/InputOps.cpp
index 3ed69a0..473e508 100644
--- a/llvm-external-projects/iree-dialects/lib/Dialect/Input/InputOps.cpp
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/Input/InputOps.cpp
@@ -8,6 +8,7 @@
 
 #include "iree-dialects/Dialect/Input/InputDialect.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -313,13 +314,57 @@
 }
 
 //===----------------------------------------------------------------------===//
+// custom<ShapedOperandList>($values, type($values), $value_dims)
+//===----------------------------------------------------------------------===//
+// %value : type{%dynamic_dims}, ...
+
+ParseResult parseShapedOperandList(
+    OpAsmParser &parser,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
+    SmallVectorImpl<Type> &valueTypes,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &valueDims) {
+  do {
+    values.emplace_back();
+    valueTypes.emplace_back();
+    if (failed(parser.parseOperand(values.back())) ||
+        failed(parser.parseColon()) ||
+        failed(parser.parseType(valueTypes.back())))
+      return failure();
+    if (int64_t dynamicDimCount =
+            cast<ShapedType>(valueTypes.back()).getNumDynamicDims()) {
+      if (failed(parser.parseOperandList(valueDims, dynamicDimCount,
+                                         AsmParser::Delimiter::Braces)))
+        return failure();
+    }
+  } while (succeeded(parser.parseOptionalComma()));
+  return success();
+}
+
+void printShapedOperandList(OpAsmPrinter &p, Operation *op, ValueRange values,
+                            TypeRange valueTypes, ValueRange valueDims) {
+  llvm::interleaveComma(llvm::zip_equal(values, valueTypes), p, [&](auto it) {
+    auto [value, valueType] = it;
+    p << value;
+    p << " : ";
+    p << valueType;
+    if (int64_t dynamicDimCount =
+            cast<ShapedType>(valueType).getNumDynamicDims()) {
+      p << "{";
+      llvm::interleaveComma(valueDims.take_front(dynamicDimCount), p);
+      valueDims = valueDims.drop_front(dynamicDimCount);
+      p << "}";
+    }
+  });
+}
+
+//===----------------------------------------------------------------------===//
 // custom<ShapedFunctionType>
 //===----------------------------------------------------------------------===//
 // (type, type{%dim0, %dim1}, type) -> (type{%dim2}, %operand4)
 
 static ParseResult
-parseShapedOperandList(OpAsmParser &parser, SmallVectorImpl<Type> &types,
-                       SmallVectorImpl<OpAsmParser::UnresolvedOperand> &dims) {
+parseShapedTypeList(OpAsmParser &parser, SmallVectorImpl<Type> &types,
+                    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &dims) {
   do {
     Type type;
     if (failed(parser.parseType(type)))
@@ -462,7 +507,7 @@
   if (failed(parser.parseLParen()))
     return failure();
   if (failed(parser.parseOptionalRParen())) {
-    if (failed(parseShapedOperandList(parser, operandTypes, operandDims)) ||
+    if (failed(parseShapedTypeList(parser, operandTypes, operandDims)) ||
         failed(parser.parseRParen())) {
       return failure();
     }
@@ -692,6 +737,25 @@
 }
 
 //===----------------------------------------------------------------------===//
+// iree_input.tensor.trace
+//===----------------------------------------------------------------------===//
+
+void TensorTraceOp::build(OpBuilder &builder, OperationState &state,
+                          StringRef key, ValueRange values) {
+  SmallVector<Value> dynamicDims;
+  for (auto value : values) {
+    auto valueType = cast<ShapedType>(value.getType());
+    for (unsigned i = 0; i < valueType.getRank(); ++i) {
+      if (valueType.isDynamicDim(i)) {
+        dynamicDims.push_back(
+            builder.createOrFold<tensor::DimOp>(state.location, value, i));
+      }
+    }
+  }
+  build(builder, state, key, values, dynamicDims);
+}
+
+//===----------------------------------------------------------------------===//
 // iree_input.dispatch
 //===----------------------------------------------------------------------===//