AFAICT costs must be constant, which limits the ability to reason about optimization.

One thing that would be interesting would be to parametrize cost based on the function inputs. For example the cost of an array elementwise addition operation could depend on the shapes of the inputs (with a default if these are not known). This could allow preferring programs were reductions are done before elementwise operations.

Interval reasoning on costs would also be nice. Even if we don't know the shape of an input, it should always be smaller after a reduction than before, and hence the cost of operating elemwise on it smaller afterwards (ignoring detail like memory layout)