Fix buck query //extension/flat_tensor: in OSS (pytorch#13484)

Author: swolchok

.ci/scripts/unittest-buck2.sh

@@ -11,9 +11,10 @@ set -eux
 # TODO: can't query //kernels/prim_ops because of non-buckified stuff in OSS.
 buck2 query "//backends/apple/... + //backends/example/... + \
 //backends/mediatek/... + //backends/transforms/... + \
-//backends/xnnpack/... + //configurations/... + //kernels/aten/... + \
-//kernels/optimized/... + //kernels/portable/... + //kernels/quantized/... + \
-//kernels/test/... + //runtime/... + //schema/... + //test/... + //util/..."
+//backends/xnnpack/... + //configurations/... + //extension/flat_tensor: + \
+//kernels/aten/... + //kernels/optimized/... + //kernels/portable/... + \
+//kernels/quantized/... + //kernels/test/... + //runtime/... + //schema/... \
++ //test/... + //util/..."
 
 # TODO: optimized ops are unbuildable because they now use ATen; put
 # them back after we can use PyTorch in OSS buck.

extension/flat_tensor/targets.bzl

@@ -1,7 +1,7 @@
-load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "get_aten_mode_options", "runtime")
 
 def define_common_targets():
-    for aten_mode in [True, False]:
+    for aten_mode in get_aten_mode_options():
         aten_suffix = "_aten" if aten_mode else ""
         runtime.cxx_library(
             name = "flat_tensor_data_map" + aten_suffix,