Repository: nod-ai/shark-ai (language: python)

Issue:
(sfnp) add permute host op.
Some use cases require a bit of array handling in between shortfin program invocations or as pre/post-processing steps. In SD workflows, it's common to see a permutation of the image output array in preparation for image library (e.g. pillow) 
consumption.

Perhaps implementing the equivalent of np.transpose ([docs](https://numpy.org/doc/stable/reference/generated/numpy.transpose.html)) would achieve what we need for most cases.

The following is a snippet showing how we use np.transpose in SDXL:
```
permuted = (
    np.transpose(vae_output, (0, 2, 3, 1))[0]
)
```



Discussion / hints:

Fix patch:
diff --git a/shortfin/python/array_binding.cc b/shortfin/python/array_binding.cc
index a05232674..08a4071a8 100644
--- a/shortfin/python/array_binding.cc
+++ b/shortfin/python/array_binding.cc
@@ -531,22 +531,26 @@ void BindArray(py::module_ &m) {
                  ->AddAsInvocationArgument(
                      inv, static_cast<local::ProgramResourceBarrier>(barrier));
            })
-      .def_static("for_device",
-                  [](local::ScopedDevice &device, std::span<const size_t> shape,
-                     DType dtype) {
-                    return custom_new_keep_alive<device_array>(
-                        py::type<device_array>(),
-                        /*keep_alive=*/device.fiber(),
-                        device_array::for_device(device, shape, dtype));
-                  })
-      .def_static("for_host",
-                  [](local::ScopedDevice &device, std::span<const size_t> shape,
-                     DType dtype) {
-                    return custom_new_keep_alive<device_array>(
-                        py::type<device_array>(),
-                        /*keep_alive=*/device.fiber(),
-                        device_array::for_host(device, shape, dtype));
-                  })
+      .def_static(
+          "for_device",
+          [](local::ScopedDevice &device, std::span<const size_t> shape,
+             DType dtype) {
+            return custom_new_keep_alive<device_array>(
+                py::type<device_array>(),
+                /*keep_alive=*/device.fiber(),
+                device_array::for_device(device, shape, dtype));
+          },
+          py::arg("device"), py::arg("shape"), py::arg("dtype"))
+      .def_static(
+          "for_host",
+          [](local::ScopedDevice &device, std::span<const size_t> shape,
+             DType dtype) {
+            return custom_new_keep_alive<device_array>(
+                py::type<device_array>(),
+                /*keep_alive=*/device.fiber(),
+                device_array::for_host(device, shape, dtype));
+          },
+          py::arg("device"), py::arg("shape"), py::arg("dtype"))
       .def("for_transfer",
            [](device_array &self) {
              return custom_new_keep_alive<device_array>(
diff --git a/shortfin/python/array_host_ops.cc b/shortfin/python/array_host_ops.cc
index 86385cfee..3e2a8ebe3 100644
--- a/shortfin/python/array_host_ops.cc
+++ b/shortfin/python/array_host_ops.cc
@@ -91,6 +91,18 @@ static const char DOCSTRING_RANDOM_GENERATOR[] =
       fixed number.
   )";
 
+static const char DOCSTRING_TRANSPOSE[] =
+    R"(Transposes axes of an array according to a permutation vector.
+
+Args:
+  input: Array to transpose.
+  permutation: New sequence of axes. Must have same number of elements as the
+    rank of input.
+  out: If given, then the results are written to this array.
+  device_visible: Whether to make the result array visible to devices. Defaults
+    to False.
+)";
+
 #define SF_UNARY_FUNCTION_CASE(dtype_name, cpp_type) \
   case DType::dtype_name():                          \
     return compute.template operator()<cpp_type>()
@@ -100,6 +112,25 @@ static const char DOCSTRING_RANDOM_GENERATOR[] =
     compute.template operator()<cpp_type>();      \
     break
 
+#define SF_MOVEMENT_OP_SWITCH(dtype)                                     \
+  if (!dtype.is_byte_aligned())                                          \
+    throw std::invalid_argument(                                         \
+        "data movement ops are only defined for byte aligned dtypes");   \
+  switch (dtype.dense_byte_count()) {                                    \
+    case 1:                                                              \
+      return compute.template operator()<uint8_t>();                     \
+    case 2:                                                              \
+      return compute.template operator()<uint16_t>();                    \
+    case 4:                                                              \
+      return compute.template operator()<uint32_t>();                    \
+    case 8:                                                              \
+      return compute.template operator()<uint64_t>();                    \
+    default:                                                             \
+      throw std::invalid_argument(                                       \
+          "data movement ops are only defined for dtypes of size 1, 2, " \
+          "4, 8");                                                       \
+  }
+
 struct PyRandomGenerator {
  public:
   using SeedType = xt::random::default_engine_type::result_type;
@@ -374,6 +405,227 @@ struct ConvertTruncFunctor {
   }
 };
 
+void OptionalArrayCast(py::handle handle,
+                       std::optional<device_array> &maybe_array) {
+  if (py::isinstance<device_array>(handle)) {
+    maybe_array.emplace(py::cast<device_array>(handle));
+  }
+}
+
+int DTypePromotionRank(DType dtype) {
+  int rank = 1;
+  if (dtype.is_boolean())
+    rank *= 1000;
+  else if (dtype.is_integer())
+    rank *= 2000;
+  else if (dtype.is_float())
+    rank *= 4000;
+  else if (dtype.is_complex())
+    rank *= 8000;
+  return rank + dtype.bit_count();
+}
+
+DType PromoteArithmeticTypes(std::optional<DType> lhs_dtype,
+                             std::optional<DType> rhs_dtype) {
+  if (!lhs_dtype && !rhs_dtype) {
+    throw std::invalid_argument(
+        "Elementwise operators require at least one argument to be a "
+        "device_array");
+  }
+
+  // One not an array: promote to the array type.
+  if (!lhs_dtype)
+    return *rhs_dtype;
+  else if (!rhs_dtype)
+    return *lhs_dtype;
+
+  int lhs_rank = DTypePromotionRank(*lhs_dtype);
+  int rhs_rank = DTypePromotionRank(*rhs_dtype);
+  DType promoted_dtype = lhs_rank < rhs_rank ? *rhs_dtype : *lhs_dtype;
+
+  // If mismatched signed/unsigned, then need to promote to the next signed
+  // dtype.
+  if (promoted_dtype.is_integer()) {
+    bool lhs_unsigned = iree_all_bits_set(
+        lhs_dtype->numerical_type(), IREE_HAL_NUMERICAL_TYPE_INTEGER_UNSIGNED);
+    bool rhs_unsigned = iree_all_bits_set(
+        rhs_dtype->numerical_type(), IREE_HAL_NUMERICAL_TYPE_INTEGER_UNSIGNED);
+    if ((lhs_unsigned || rhs_unsigned) && !(lhs_unsigned && rhs_unsigned)) {
+      // Signed/unsigned mismatch. Promote to next.
+      switch (promoted_dtype) {
+        case DType::uint8():
+        case DType::int8():
+          return DType::int16();
+        case DType::uint16():
+        case DType::int16():
+          return DType::int32();
+        case DType::uint32():
+        case DType::int32():
+          return DType::int64();
+        default:
+          // Jax's type promotion chart says this goes to a weak FP type, but
+          // we don't implement such a construct and I don't really see how
+          // that makes sense in a system setting like this, so we just saturate
+          // to 64bit.
+          return DType::int64();
+      }
+    }
+  }
+
+  return promoted_dtype;
+}
+
+// ---------------------------------------------------------------------------//
+// Elementwise support
+// ---------------------------------------------------------------------------//
+
+// Python element type scalar conversion functions.
+uint8_t ConvertPyToEltTy(py::handle py_value, uint8_t zero) {
+  return py::cast<uint8_t>(py_value);
+}
+
+int8_t ConvertPyToEltTy(py::handle py_value, int8_t zero) {
+  return py::cast<int8_t>(py_value);
+}
+
+uint16_t ConvertPyToEltTy(py::handle py_value, uint16_t zero) {
+  return py::cast<uint16_t>(py_value);
+}
+
+int16_t ConvertPyToEltTy(py::handle py_value, int16_t zero) {
+  return py::cast<int16_t>(py_value);
+}
+
+uint32_t ConvertPyToEltTy(py::handle py_value, uint32_t zero) {
+  return py::cast<uint32_t>(py_value);
+}
+
+int32_t ConvertPyToEltTy(py::handle py_value, int32_t zero) {
+  return py::cast<int32_t>(py_value);
+}
+
+uint64_t ConvertPyToEltTy(py::handle py_value, uint64_t zero) {
+  return py::cast<uint64_t>(py_value);
+}
+
+int64_t ConvertPyToEltTy(py::handle py_value, int64_t zero) {
+  return py::cast<int64_t>(py_value);
+}
+
+float ConvertPyToEltTy(py::handle py_value, float zero) {
+  return py::cast<float>(py_value);
+}
+
+double ConvertPyToEltTy(py::handle py_value, double zero) {
+  return py::cast<double>(py_value);
+}
+
+half_float::half ConvertPyToEltTy(py::handle py_value, half_float::half zero) {
+  // Python can't cast directly to half so first go to double.
+  return static_cast<half_float::half>(py::cast<double>(py_value));
+}
+
+struct AddFunctor {
+  template <typename Lhs, typename Rhs>
+  static auto Invoke(Lhs &&lhs, Rhs &&rhs) {
+    return lhs + rhs;
+  }
+};
+
+struct DivideFunctor {
+  template <typename Lhs, typename Rhs>
+  static auto Invoke(Lhs &&lhs, Rhs &&rhs) {
+    return lhs / rhs;
+  }
+};
+
+struct MultiplyFunctor {
+  template <typename Lhs, typename Rhs>
+  static auto Invoke(Lhs &&lhs, Rhs &&rhs) {
+    return lhs * rhs;
+  }
+};
+
+struct SubtractFunctor {
+  template <typename Lhs, typename Rhs>
+  static auto Invoke(Lhs &&lhs, Rhs &&rhs) {
+    return lhs - rhs;
+  }
+};
+
+template <typename ElementwiseFunctor>
+device_array ElementwiseOperation(py::handle lhs, py::handle rhs,
+                                  std::optional<device_array> out,
+                                  bool device_visible) {
+  std::optional<device_array> lhs_array;
+  OptionalArrayCast(lhs, lhs_array);
+  std::optional<device_array> rhs_array;
+  OptionalArrayCast(rhs, rhs_array);
+  auto dtype = PromoteArithmeticTypes(
+      lhs_array ? std::optional<DType>(lhs_array->dtype()) : std::nullopt,
+      rhs_array ? std::optional<DType>(rhs_array->dtype()) : std::nullopt);
+  if (lhs_array && lhs_array->dtype() != dtype) {
+    auto converted = GenericElementwiseConvert<ConvertFunctor>(
+        *lhs_array, dtype, /*out=*/std::nullopt,
+        /*device_visible=*/false);
+    lhs_array.reset();
+    lhs_array.emplace(std::move(converted));
+  }
+  if (rhs_array && rhs_array->dtype() != dtype) {
+    auto converted = GenericElementwiseConvert<ConvertFunctor>(
+        *rhs_array, dtype, /*out=*/std::nullopt,
+        /*device_visible=*/false);
+    rhs_array.reset();
+    rhs_array.emplace(std::move(converted));
+  }
+
+  auto compute = [&]<typename EltTy>() -> device_array {
+    auto handle_result = [&]<typename D, typename A>(
+                             D &&device, A &&result) -> device_array {
+      if (!out) {
+        out.emplace(device_array::for_host(device, result.shape(), dtype,
+                                           device_visible));
+      }
+      auto out_t = out->map_xtensor_w<EltTy>();
+      *out_t = result;
+      return *out;
+    };
+    if (!rhs_array) {
+      auto lhs_t = lhs_array->map_xtensor<EltTy>();
+      xt::xarray<EltTy> rhs_scalar = ConvertPyToEltTy(rhs, EltTy());
+      return handle_result(lhs_array->device(),
+                           ElementwiseFunctor::Invoke(*lhs_t, rhs_scalar));
+    } else if (!lhs_array) {
+      xt::xarray<EltTy> lhs_scalar = ConvertPyToEltTy(lhs, EltTy());
+      auto rhs_t = rhs_array->map_xtensor<EltTy>();
+      return handle_result(rhs_array->device(),
+                           ElementwiseFunctor::Invoke(lhs_scalar, *rhs_t));
+    } else {
+      auto lhs_t = lhs_array->map_xtensor<EltTy>();
+      auto rhs_t = rhs_array->map_xtensor<EltTy>();
+      return handle_result(lhs_array->device(),
+                           ElementwiseFunctor::Invoke(*lhs_t, *rhs_t));
+    }
+  };
+
+  switch (dtype) {
+    SF_UNARY_FUNCTION_CASE(float16, half_float::half);
+    SF_UNARY_FUNCTION_CASE(float32, float);
+    SF_UNARY_FUNCTION_CASE(float64, double);
+    SF_UNARY_FUNCTION_CASE(uint8, uint8_t);
+    SF_UNARY_FUNCTION_CASE(int8, int8_t);
+    SF_UNARY_FUNCTION_CASE(uint16, uint16_t);
+    SF_UNARY_FUNCTION_CASE(int16, int16_t);
+    SF_UNARY_FUNCTION_CASE(uint32, uint32_t);
+    SF_UNARY_FUNCTION_CASE(int32, uint32_t);
+    SF_UNARY_FUNCTION_CASE(uint64, uint64_t);
+    SF_UNARY_FUNCTION_CASE(int64, int64_t);
+    default:
+      throw std::invalid_argument(fmt::format(
+          "Unsupported dtype({}) for in elementwise op", dtype.name()));
+  }
+}
+
 }  // namespace
 
 void BindArrayHostOps(py::module_ &m) {
@@ -457,6 +709,39 @@ void BindArrayHostOps(py::module_ &m) {
   SF_DEF_CONVERT("floor", GenericElementwiseConvert<ConvertFloorFunctor>);
   SF_DEF_CONVERT("round", GenericElementwiseConvert<ConvertRoundFunctor>);
   SF_DEF_CONVERT("trunc", GenericElementwiseConvert<ConvertTruncFunctor>);
-}
+
+  // Transpose.
+  m.def(
+      "transpose",
+      [](device_array input, std::vector<size_t> permutation,
+         std::optional<device_array> out, bool device_visible) {
+        auto compute = [&]<typename EltTy>() -> device_array {
+          auto input_t = input.map_xtensor<EltTy>();
+          auto permuted_t =
+              xt::transpose(*input_t, permutation, xt::check_policy::full());
+          if (!out) {
+            out.emplace(device_array::for_host(input.device(),
+                                               permuted_t.shape(),
+                                               input.dtype(), device_visible));
+          }
+          auto out_t = out->map_xtensor_w<EltTy>();
+          *out_t = permuted_t;
+          return *out;
+        };
+        SF_MOVEMENT_OP_SWITCH(input.dtype());
+      },
+      py::arg("input"), py::arg("permutation"), py::arg("out") = py::none(),
+      py::arg("device_visible") = false, DOCSTRING_TRANSPOSE);
+
+// Elementwise.
+#define SF_DEF_ELEMENTWISE(py_name, target)                             \
+  m.def(py_name, target, py::arg("lhs"), py::arg("rhs"), py::kw_only(), \
+        py::arg("out") = py::none(), py::arg("device_visible") = false)
+  SF_DEF_ELEMENTWISE("add", ElementwiseOperation<AddFunctor>);
+  SF_DEF_ELEMENTWISE("divide", ElementwiseOperation<DivideFunctor>);
+  SF_DEF_ELEMENTWISE("multiply", ElementwiseOperation<MultiplyFunctor>);
+  SF_DEF_ELEMENTWISE("subtract", ElementwiseOperation<SubtractFunctor>);
+
+}  // namespace shortfin::python
 
 }  // namespace shortfin::python
diff --git a/shortfin/python/shortfin/array/__init__.py b/shortfin/python/shortfin/array/__init__.py
index 6079541c8..670102dfe 100644
--- a/shortfin/python/shortfin/array/__init__.py
+++ b/shortfin/python/shortfin/array/__init__.py
@@ -44,11 +44,16 @@
 
 # Ops.
 argmax = _sfl.array.argmax
+add = _sfl.array.add
 ceil = _sfl.array.ceil
 convert = _sfl.array.convert
+divide = _sfl.array.divide
 fill_randn = _sfl.array.fill_randn
 floor = _sfl.array.floor
+multiply = _sfl.array.multiply
 round = _sfl.array.round
+subtract = _sfl.array.subtract
+transpose = _sfl.array.transpose
 trunc = _sfl.array.trunc
 RandomGenerator = _sfl.array.RandomGenerator
 
@@ -86,12 +91,17 @@
     "storage",
     "DType",
     # Ops.
+    "add",
     "argmax",
     "ceil",
     "convert",
+    "divide",
     "fill_randn",
     "floor",
+    "multiply",
     "round",
+    "subtract",
+    "transpose",
     "trunc",
     "RandomGenerator",
 ]
diff --git a/shortfin/src/shortfin/array/dtype.h b/shortfin/src/shortfin/array/dtype.h
index d746d69bf..de1763698 100644
--- a/shortfin/src/shortfin/array/dtype.h
+++ b/shortfin/src/shortfin/array/dtype.h
@@ -49,6 +49,9 @@ class SHORTFIN_API DType {
   bool is_integer_bitwidth(size_t bitwidth) const {
     return iree_hal_element_type_is_integer(et_, bitwidth);
   }
+  uint32_t numerical_type() const {
+    return iree_hal_element_numerical_type(et_);
+  }
 
   // Computes the size in bytes required to store densely packed nd-dims.
   // This presently only supports byte aligned dtypes. In the future, when


Test patch:
diff --git a/shortfin/tests/api/array_ops_test.py b/shortfin/tests/api/array_ops_test.py
index 7c792d92b..164dfb479 100644
--- a/shortfin/tests/api/array_ops_test.py
+++ b/shortfin/tests/api/array_ops_test.py
@@ -268,3 +268,171 @@ def test_nearest_int_conversion(device, dtype, out_dtype, sfnp_func, ref_round_f
     assert output.dtype == out_dtype
     for ref, actual in zip(ref_rounded, output.items):
         assert ref == int(actual)
+
+
+def test_elementwise_forms(device):
+    # All elementwise ops use the same template expansion which enforces
+    # certain common invariants. Here we test these on the multiply op,
+    # relying on a parametric test for actual behavior.
+    with pytest.raises(
+        ValueError,
+        match="Elementwise operators require at least one argument to be a device_array",
+    ):
+        sfnp.multiply(2, 2)
+
+    ary = sfnp.device_array.for_host(device, [2, 3], dtype=sfnp.float32)
+    with ary.map(discard=True) as m:
+        m.fill(42.0)
+
+    # Rhs scalar int accepted.
+    result = sfnp.multiply(ary, 2)
+    assert list(result.items) == [84.0] * 6
+
+    # Rhs scalar float accepted.
+    result = sfnp.multiply(ary, 2.0)
+    assert list(result.items) == [84.0] * 6
+
+    # Lhs scalar int accepted.
+    result = sfnp.multiply(2, ary)
+    assert list(result.items) == [84.0] * 6
+
+    # Lhs scalar float accepted.
+    result = sfnp.multiply(2.0, ary)
+    assert list(result.items) == [84.0] * 6
+
+    # Out.
+    out = sfnp.device_array.for_host(device, [2, 3], dtype=sfnp.float32)
+    sfnp.multiply(2.0, ary, out=out)
+    assert list(out.items) == [84.0] * 6
+
+
+@pytest.mark.parametrize(
+    "lhs_dtype,rhs_dtype,promoted_dtype",
+    [
+        (sfnp.float32, sfnp.float16, sfnp.float32),
+        (sfnp.float16, sfnp.float32, sfnp.float32),
+        (sfnp.float32, sfnp.float64, sfnp.float64),
+        (sfnp.float64, sfnp.float32, sfnp.float64),
+        # Integer promotion.
+        (sfnp.uint8, sfnp.uint16, sfnp.uint16),
+        (sfnp.uint16, sfnp.uint32, sfnp.uint32),
+        (sfnp.uint32, sfnp.uint64, sfnp.uint64),
+        (sfnp.int8, sfnp.int16, sfnp.int16),
+        (sfnp.int16, sfnp.int32, sfnp.int32),
+        (sfnp.int32, sfnp.int64, sfnp.int64),
+        # Signed/unsigned promotion.
+        (sfnp.int8, sfnp.uint8, sfnp.int16),
+        (sfnp.int16, sfnp.uint16, sfnp.int32),
+        (sfnp.int32, sfnp.uint32, sfnp.int64),
+        (sfnp.int8, sfnp.uint32, sfnp.int64),
+    ],
+)
+def test_elementwise_promotion(device, lhs_dtype, rhs_dtype, promoted_dtype):
+    # Tests that promotion infers an appropriate result type.
+    lhs = sfnp.device_array.for_host(device, [2, 3], lhs_dtype)
+    rhs = sfnp.device_array.for_host(device, [2, 3], rhs_dtype)
+    result = sfnp.multiply(lhs, rhs)
+    assert result.dtype == promoted_dtype
+
+
+@pytest.mark.parametrize(
+    "dtype,op,check_value",
+    [
+        # Add.
+        (sfnp.int8, sfnp.add, 44.0),
+        (sfnp.int16, sfnp.add, 44.0),
+        (sfnp.int32, sfnp.add, 44.0),
+        (sfnp.int64, sfnp.add, 44.0),
+        (sfnp.uint8, sfnp.add, 44.0),
+        (sfnp.uint16, sfnp.add, 44.0),
+        (sfnp.uint32, sfnp.add, 44.0),
+        (sfnp.uint64, sfnp.add, 44.0),
+        (sfnp.float16, sfnp.add, 44.0),
+        (sfnp.float32, sfnp.add, 44.0),
+        (sfnp.float64, sfnp.add, 44.0),
+        # Divide.
+        (sfnp.int8, sfnp.divide, 21.0),
+        (sfnp.int16, sfnp.divide, 21.0),
+        (sfnp.int32, sfnp.divide, 21.0),
+        (sfnp.int64, sfnp.divide, 21.0),
+        (sfnp.uint8, sfnp.divide, 21.0),
+        (sfnp.uint16, sfnp.divide, 21.0),
+        (sfnp.uint32, sfnp.divide, 21.0),
+        (sfnp.uint64, sfnp.divide, 21.0),
+        (sfnp.float16, sfnp.divide, 21.0),
+        (sfnp.float32, sfnp.divide, 21.0),
+        (sfnp.float64, sfnp.divide, 21.0),
+        # Multiply.
+        (sfnp.int8, sfnp.multiply, 84.0),
+        (sfnp.int16, sfnp.multiply, 84.0),
+        (sfnp.int32, sfnp.multiply, 84.0),
+        (sfnp.int64, sfnp.multiply, 84.0),
+        (sfnp.uint8, sfnp.multiply, 84.0),
+        (sfnp.uint16, sfnp.multiply, 84.0),
+        (sfnp.uint32, sfnp.multiply, 84.0),
+        (sfnp.uint64, sfnp.multiply, 84.0),
+        (sfnp.float16, sfnp.multiply, 84.0),
+        (sfnp.float32, sfnp.multiply, 84.0),
+        (sfnp.float64, sfnp.multiply, 84.0),
+        # Subtract.
+        (sfnp.int8, sfnp.subtract, 40.0),
+        (sfnp.int16, sfnp.subtract, 40.0),
+        (sfnp.int32, sfnp.subtract, 40.0),
+        (sfnp.int64, sfnp.subtract, 40.0),
+        (sfnp.uint8, sfnp.subtract, 40.0),
+        (sfnp.uint16, sfnp.subtract, 40.0),
+        (sfnp.uint32, sfnp.subtract, 40.0),
+        (sfnp.uint64, sfnp.subtract, 40.0),
+        (sfnp.float16, sfnp.subtract, 40.0),
+        (sfnp.float32, sfnp.subtract, 40.0),
+        (sfnp.float64, sfnp.subtract, 40.0),
+    ],
+)
+def test_elementwise_array_correctness(device, dtype, op, check_value):
+    lhs = sfnp.device_array.for_host(device, [2, 2], sfnp.int32)
+    with lhs.map(discard=True) as m:
+        m.fill(42)
+
+    rhs = sfnp.device_array.for_host(device, [2], sfnp.int32)
+    with rhs.map(discard=True) as m:
+        m.fill(2)
+
+    lhs = sfnp.convert(lhs, dtype=dtype)
+    rhs = sfnp.convert(rhs, dtype=dtype)
+    result = op(lhs, rhs)
+    assert result.shape == [2, 2]
+    result = sfnp.convert(result, dtype=sfnp.float32)
+    items = list(result.items)
+    assert items == [check_value] * 4
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        sfnp.int8,
+        sfnp.int16,
+        sfnp.int32,
+        sfnp.int64,
+        sfnp.uint8,
+        sfnp.uint16,
+        sfnp.uint32,
+        sfnp.uint64,
+        sfnp.float32,
+        sfnp.float16,
+        sfnp.float32,
+        sfnp.float64,
+    ],
+)
+def test_transpose(device, dtype):
+    input = sfnp.device_array.for_host(device, [3, 2], sfnp.int32)
+    input.items = [0, 1, 2, 3, 4, 5]
+    input = sfnp.convert(input, dtype=dtype)
+    permuted = sfnp.transpose(input, [1, 0])
+    assert permuted.shape == [2, 3]
+    items = list(sfnp.convert(permuted, dtype=sfnp.int32).items)
+    assert items == [0, 2, 4, 1, 3, 5]
+
+    out = sfnp.device_array.for_host(device, [2, 3], dtype)
+    sfnp.transpose(input, [1, 0], out=out)
+    items = list(sfnp.convert(permuted, dtype=sfnp.int32).items)
+    assert items == [0, 2, 4, 1, 3, 5]
diff --git a/shortfin/tests/api/array_use_case_test.py b/shortfin/tests/api/array_use_case_test.py
new file mode 100644
index 000000000..d4a030d45
--- /dev/null
+++ b/shortfin/tests/api/array_use_case_test.py
@@ -0,0 +1,64 @@
+# Copyright 2024 Advanced Micro Devices, Inc.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+import array
+import math
+import pytest
+
+import shortfin as sf
+import shortfin.array as sfnp
+
+
+@pytest.fixture
+def lsys():
+    # TODO: Port this test to use memory type independent access. It currently
+    # presumes unified memory.
+    # sc = sf.SystemBuilder()
+    sc = sf.host.CPUSystemBuilder()
+    lsys = sc.create_system()
+    yield lsys
+    lsys.shutdown()
+
+
+@pytest.fixture
+def fiber(lsys):
+    return lsys.create_fiber()
+
+
+@pytest.fixture
+def device(fiber):
+    return fiber.device(0)
+
+
+# Tests a typical image conversion from a model oriented layout to an array
+# of contained images.
+def test_image_to_bytes(device):
+    bs = 2
+    height = 16
+    width = 12
+    images_shape = [bs, 3, height, width]
+    images_planar = sfnp.device_array.for_host(device, images_shape, sfnp.float32)
+    # Band the data so that each channel increases by 0.1 across images.
+    for i in range(bs):
+        for j in range(3):
+            data = [i * 0.3 + j * 0.1 for _ in range(height * width)]
+            images_planar.view(i, j).items = data
+    images_planar = sfnp.convert(images_planar, dtype=sfnp.float16)
+
+    # Extract and convert each image to interleaved RGB bytes.
+    images = []
+    for idx in range(images_planar.shape[0]):
+        image_planar = images_planar.view(idx)
+        assert image_planar.shape == [1, 3, 16, 12]
+        image_interleaved = sfnp.transpose(image_planar, (0, 2, 3, 1))
+        assert image_interleaved.shape == [1, 16, 12, 3]
+        image_scaled = sfnp.multiply(image_interleaved, 255)
+        image = sfnp.round(image_scaled, dtype=sfnp.uint8)
+        image_bytes = bytes(image.map(read=True))
+        images.append(image_bytes)
+
+    assert images[0] == b"\x00\x1a3" * 192
+    assert images[1] == b"Mf\x80" * 192