def all_arrays(cond=lambda x: True): a = list(ek.scalar.__dict__.items()) a += ek.packet.__dict__.items() a += ek.cuda.__dict__.items() a += ek.llvm.__dict__.items() a += ek.llvm.__dict__.items() return [v for k, v in a if isinstance(v, type) and cond(v) and not ek.is_special_v(v) and not ek.array_depth_v(v) >= 3 and not (ek.array_depth_v(v) >= 2 and 'scalar' in v.__module__)]
def full_(cls, value, size, eval): result = cls() if cls.Size == Dynamic: result.init_(size) for i in range(size): result.set_entry_(i, value) else: if _ek.array_depth_v(value) != cls.Depth - 1: value = _ek.full(cls.Value, value, size, eval) for i in range(cls.Size): result.set_entry_(i, value) return result
def mul_(a0, a1): if not a0.IsArithmetic: raise Exception("mul(): requires arithmetic operands!") if _ek.array_depth_v(a1) < _ek.array_depth_v(a0): # Avoid type promotion in scalars multiplication, which would # be costly for special types (matrices, quaternions, etc.) sr = len(a0) ar = a0.empty_(sr if a0.Size == Dynamic else 0) for i in range(len(a0)): ar[i] = a0[i] * a1 return ar ar, sr = _check2(a0, a1) if not a0.IsSpecial: for i in range(sr): ar[i] = a0[i] * a1[i] elif a0.IsComplex: ar.real = _ek.fmsub(a0.real, a1.real, a0.imag * a1.imag) ar.imag = _ek.fmadd(a0.real, a1.imag, a0.imag * a1.real) elif a0.IsQuaternion: tbl = (4, 3, -2, 1, -3, 4, 1, 2, 2, -1, 4, 3, -1, -2, -3, 4) for i in range(4): accum = 0 for j in range(4): idx = tbl[i * 4 + j] value = a1[abs(idx) - 1] accum = _ek.fmadd(a0[j], value if idx > 0 else -value, accum) ar[i] = accum elif a0.IsMatrix: raise Exception( "mul(): please use the matrix multiplication operator '@' instead." ) else: raise Exception("mul(): unsupported array type!") return ar