ak.flatten on dask_awkward BitMaskedArray triggers awkward error

[NJManganelli]

Version of Awkward Array

2.6.1

Description and code to reproduce

This only occurs with dask_awkward input arrays, but it triggers an error within awkward (let me know if I should re-file in dask_awkward instead)

File submitted for another issue and welcome to be vendored as requested, here's a link to it again:
https://drive.google.com/file/d/1MFpGxT1pWqF0cZ0VMnuszl2PXK6uBmL3/view?usp=sharing

Reproducer tested with dask_awkward versions '2024.2.1.dev16+g379d0d4' and '2024.1.2')

import awkward as ak
import dask_awkward as dak
esrc = ak.from_parquet("join_out_flat-part2.parquet")
dsrc = dak.from_parquet("join_out_flat-part2.parquet")

#awkward OKAY
ak.flatten(esrc.electron_pt, axis=None)
ak.flatten(esrc.electron_pt, axis=0)
ak.flatten(esrc.electron_pt, axis=1)
ak.flatten(esrc.electron_pt, axis=-1)

#dask_awkward array OKAY
ak.flatten(dsrc.electron_pt, axis=0).compute()
ak.flatten(dsrc.electron_pt, axis=None).compute()

ak.flatten(dsrc.electron_pt.compute(), axis=0)
ak.flatten(dsrc.electron_pt.compute(), axis=None)

#dask_awkward array ERROR in awkward
ak.flatten(dsrc.electron_pt, axis=1) #errors before .compute()
ak.flatten(dsrc.electron_pt, axis=-1) #errors before .compute()
ak.flatten(dsrc.electron_pt.compute(), axis=1)

Error triggered on uncomputed array:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[1], line 4
      2 import dask_awkward as dak
      3 dsrc = dak.from_parquet("join_out_flat-part2.parquet")
----> 4 ak.flatten(dsrc.electron_pt, axis=1)

File [/usr/local/lib/python3.11/site-packages/awkward/_dispatch.py:51](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/_dispatch.py#line=50), in named_high_level_function.<locals>.dispatch(*args, **kwargs)
     49     continue
     50 else:
---> 51     result = custom_impl(dispatch, array_likes, args, kwargs)
     53     # Future proof the implementation by permitting the `__awkward_function__` to return `NotImplemented`
     54     # This may later be used to signal that another overload should be used.
     55     if result is NotImplemented:

File /srv/.env/lib/python3.11/site-packages/dask_awkward/lib/core.py:1601, in Array.__awkward_function__(self, func, array_likes, args, kwargs)
   1599     except AttributeError:
   1600         return NotImplemented
-> 1601 return fn(*args, **kwargs)

File /srv/.env/lib/python3.11/site-packages/dask_awkward/lib/structure.py:442, in flatten(array, axis, highlevel, behavior, attrs)
    440 if not highlevel:
    441     raise ValueError("Only highlevel=True is supported")
--> 442 return map_partitions(
    443     _FlattenFn(axis=axis, highlevel=highlevel, behavior=behavior, attrs=attrs),
    444     array,
    445     label="flatten",
    446     output_divisions=None,
    447 )

File /srv/.env/lib/python3.11/site-packages/dask_awkward/lib/core.py:2029, in map_partitions(base_fn, label, token, meta, output_divisions, traverse, *args, **kwargs)
   2021 lay = partitionwise_layer(
   2022     fn,
   2023     name,
   2024     *arg_flat_deps_expanded,
   2025     *kwarg_flat_deps,
   2026 )
   2028 if meta is None:
-> 2029     meta = map_meta(fn, *arg_flat_deps_expanded, *kwarg_flat_deps)
   2031 hlg = HighLevelGraph.from_collections(
   2032     name,
   2033     lay,
   2034     dependencies=flat_deps,
   2035 )
   2037 dak_arrays = tuple(filter(lambda x: isinstance(x, Array), flat_deps))

File /srv/.env/lib/python3.11/site-packages/dask_awkward/lib/core.py:2415, in map_meta(fn, *deps)
   2411 def map_meta(fn: ArgsKwargsPackedFunction, *deps: Any) -> ak.Array | None:
   2412     # NOTE: fn is assumed to be a *packed* function
   2413     #       as defined up in map_partitions. be careful!
   2414     try:
-> 2415         meta = fn(*to_meta(deps))
   2416         return meta
   2417     except Exception as err:
   2418         # if compute-unknown-meta is False then we don't care about
   2419         # this failure and we return None.

File /srv/.env/lib/python3.11/site-packages/dask_awkward/lib/core.py:1890, in ArgsKwargsPackedFunction.__call__(self, *args_deps_expanded)
   1888     len_args += n_args
   1889 kwargs = self.kwarg_repacker(args_deps_expanded[len_args:])[0]
-> 1890 return self.fn(*args, **kwargs)

File /srv/.env/lib/python3.11/site-packages/dask_awkward/lib/structure.py:429, in _FlattenFn.__call__(self, array)
    428 def __call__(self, array: ak.Array) -> ak.Array:
--> 429     return ak.flatten(array, **self.kwargs)

File [/usr/local/lib/python3.11/site-packages/awkward/_dispatch.py:62](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/_dispatch.py#line=61), in named_high_level_function.<locals>.dispatch(*args, **kwargs)
     60 # Failed to find a custom overload, so resume the original function
     61 try:
---> 62     next(gen_or_result)
     63 except StopIteration as err:
     64     return err.value

File [/usr/local/lib/python3.11/site-packages/awkward/operations/ak_flatten.py:172](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/operations/ak_flatten.py#line=171), in flatten(array, axis, highlevel, behavior, attrs)
    169 yield (array,)
    171 # Implementation
--> 172 return _impl(array, axis, highlevel, behavior, attrs)

File [/usr/local/lib/python3.11/site-packages/awkward/operations/ak_flatten.py:236](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/operations/ak_flatten.py#line=235), in _impl(array, axis, highlevel, behavior, attrs)
    234     out = apply(layout)
    235 else:
--> 236     out = ak._do.flatten(layout, axis)
    237 return ctx.wrap(out, highlevel=highlevel)

File [/usr/local/lib/python3.11/site-packages/awkward/_do.py:233](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/_do.py#line=232), in flatten(layout, axis)
    232 def flatten(layout: Content, axis: int = 1) -> Content:
--> 233     offsets, flattened = layout._offsets_and_flattened(axis, 1)
    234     return flattened

File [/usr/local/lib/python3.11/site-packages/awkward/contents/bitmaskedarray.py:571](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/contents/bitmaskedarray.py#line=570), in BitMaskedArray._offsets_and_flattened(self, axis, depth)
    570 def _offsets_and_flattened(self, axis: int, depth: int) -> tuple[Index, Content]:
--> 571     return self.to_ByteMaskedArray._offsets_and_flattened(axis, depth)

AttributeError: 'function' object has no attribute '_offsets_and_flattened'

This error occurred while calling

    ak.flatten(
        Array-instance
        axis = 1
    )

Error triggered on already-computed array:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[1], line 4
      2 import dask_awkward as dak
      3 src = dak.from_parquet("join_out_flat-part2.parquet")
----> 4 ak.flatten(src.electron_pt.compute(), axis=1)

File [/usr/local/lib/python3.11/site-packages/awkward/_dispatch.py:62](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/_dispatch.py#line=61), in named_high_level_function.<locals>.dispatch(*args, **kwargs)
     60 # Failed to find a custom overload, so resume the original function
     61 try:
---> 62     next(gen_or_result)
     63 except StopIteration as err:
     64     return err.value

File [/usr/local/lib/python3.11/site-packages/awkward/operations/ak_flatten.py:172](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/operations/ak_flatten.py#line=171), in flatten(array, axis, highlevel, behavior, attrs)
    169 yield (array,)
    171 # Implementation
--> 172 return _impl(array, axis, highlevel, behavior, attrs)

File [/usr/local/lib/python3.11/site-packages/awkward/operations/ak_flatten.py:236](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/operations/ak_flatten.py#line=235), in _impl(array, axis, highlevel, behavior, attrs)
    234     out = apply(layout)
    235 else:
--> 236     out = ak._do.flatten(layout, axis)
    237 return ctx.wrap(out, highlevel=highlevel)

File [/usr/local/lib/python3.11/site-packages/awkward/_do.py:233](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/_do.py#line=232), in flatten(layout, axis)
    232 def flatten(layout: Content, axis: int = 1) -> Content:
--> 233     offsets, flattened = layout._offsets_and_flattened(axis, 1)
    234     return flattened

File [/usr/local/lib/python3.11/site-packages/awkward/contents/bitmaskedarray.py:571](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/contents/bitmaskedarray.py#line=570), in BitMaskedArray._offsets_and_flattened(self, axis, depth)
    570 def _offsets_and_flattened(self, axis: int, depth: int) -> tuple[Index, Content]:
--> 571     return self.to_ByteMaskedArray._offsets_and_flattened(axis, depth)

AttributeError: 'function' object has no attribute '_offsets_and_flattened'

This error occurred while calling

    ak.flatten(
        <Array [[], [24.9, ...], ..., [59.3, 5.22]] type='25763 * option[va...'>
        axis = 1
    )

type of post-computed array:

#post
type: 25763 * option[var * ?float32]
[[],
 [24.9, 6.66],
 [25.4],
 [46.3, 14.1, 10.6],
 [11.8, 9.79],
 [],
 [6.71],
 [38.7, 36.7, 23.9, 12.7],
 [44.7],
 [38.4],
 ...,
 [22.4, 12.9],
 [28.9],
 [5.18],
 [],
 [],
 [20.5],
 [25.3, 10],
 [8.42],
 [59.3, 5.22]]