Added opaque_size_ functions

Author: DoeringChristian

ext/drjit-core

@@ -526,9 +526,13 @@ DRJIT_ROUTE_BINARY_FALLBACK(dot, dot, (E) a1 * (E) a2)
 
 template <typename Array>
 DRJIT_INLINE auto mean(const Array &a) {
-    if constexpr (is_array_v<Array>)
-        return sum(a) * (1.f / a.derived().size());
-    else
+    if constexpr (is_array_v<Array>){
+        if (jit_flag(JitFlag::FreezingScope)) {
+            return sum(a) * (1.f / a.derived().opaque_size_());
+        } else {
+            return sum(a) * (1.f / a.derived().size());
+        }
+    } else
         return a;
 }
 

include/drjit/array_router.h

@@ -79,6 +79,8 @@ struct StaticArrayBase : ArrayBaseT<Value_, IsMask_, Derived_> {
 
     DRJIT_INLINE constexpr size_t size() const { return Derived::Size; }
 
+    DRJIT_INLINE constexpr size_t opaque_size_() const { return Derived::Size; }
+
     DRJIT_INLINE void init_(size_t) { }
 
     static Derived empty_(size_t size) {

include/drjit/array_static.h

@@ -709,6 +709,11 @@ struct DRJIT_TRIVIAL_ABI DiffArray
 
     size_t size() const { return jit_var_size((uint32_t) m_index); }
 
+    auto opaque_size_() const {
+        using UInt32 = JitArray<Backend, uint32_t>;
+        return UInt32::steal(jit_var_opaque_width(m_index));
+    }
+
     bool grad_enabled_() const {
         if constexpr (IsFloat)
             return (m_index >> 32) != 0;

include/drjit/autodiff.h

@@ -123,6 +123,7 @@ struct DynamicArray
     }
 
     DRJIT_INLINE size_t size() const { return m_size; }
+    DRJIT_INLINE size_t opaque_size_() const { return m_size; }
     DRJIT_INLINE DynamicArray copy() { return DynamicArray(*this); }
 
     DRJIT_INLINE Value &entry(size_t i) {

include/drjit/dynamic.h

@@ -596,6 +596,10 @@ struct DRJIT_TRIVIAL_ABI JitArray
 
     bool valid() const { return m_index != 0; }
     size_t size() const { return jit_var_size(m_index); }
+    auto opaque_size_() const {
+        using UInt32 = JitArray<Backend, uint32_t>;
+        return UInt32::steal(jit_var_opaque_width(m_index));
+    }
     uint32_t index() const { return m_index; }
     uint32_t index_ad() const { return 0; }
     uint64_t index_combined() const { return m_index; }